# # SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import os import torch from demo_diffusion.dynamic_import import import_from_diffusers from demo_diffusion.model import base_model, load, optimizer # List of models to import from diffusers.models models_to_import = ["SD3Transformer2DModel", "SD3ControlNetModel"] for model in models_to_import: globals()[model] = import_from_diffusers(model, "diffusers.models") class SD3ControlNetWrapper(torch.nn.Module): def __init__(self, controlnet): super().__init__() self.controlnet = controlnet def forward(self, hidden_states, controlnet_cond, conditioning_scale, pooled_projections, timestep): params = { "hidden_states": hidden_states, "pooled_projections": pooled_projections, "timestep": timestep, "controlnet_cond": controlnet_cond, "conditioning_scale": conditioning_scale, } out = self.controlnet(**params)["controlnet_block_samples"] return torch.stack(out, dim=0) class SD3ControlNet(base_model.BaseModel): def __init__( self, version, controlnet, pipeline, device, hf_token, verbose, framework_model_dir, fp16=False, tf32=False, bf16=False, int8=False, fp8=False, max_batch_size=16, do_classifier_free_guidance=False, ): super(SD3ControlNet, self).__init__( version, pipeline, device=device, hf_token=hf_token, verbose=verbose, framework_model_dir=framework_model_dir, fp16=fp16, tf32=tf32, bf16=bf16, int8=int8, fp8=fp8, max_batch_size=max_batch_size, ) self.path = load.get_path(version, pipeline, controlnet) self.subfolder = "controlnet_{}".format(controlnet) self.controlnet_model_dir = load.get_checkpoint_dir( self.framework_model_dir, self.version, self.pipeline, self.subfolder ) self.transformer_model_dir = load.get_checkpoint_dir( self.framework_model_dir, self.version, self.pipeline, "transformer" ) if not os.path.exists(self.controlnet_model_dir): self.config = SD3ControlNetModel.load_config(self.path, token=self.hf_token) else: print(f"[I] Load SD3ControlNetModel config from: {self.controlnet_model_dir}") self.config = SD3ControlNetModel.load_config(self.controlnet_model_dir) self.xB = 2 if do_classifier_free_guidance else 1 # batch multiplier def get_model(self, torch_inference=""): model_opts = ( {"torch_dtype": torch.float16} if self.fp16 else {"torch_dtype": torch.bfloat16} if self.bf16 else {} ) if not load.is_model_cached(self.controlnet_model_dir, model_opts, self.hf_safetensor): model = SD3ControlNetModel.from_pretrained(self.path, **model_opts, use_safetensors=self.hf_safetensor).to( self.device ) model.save_pretrained(self.controlnet_model_dir, **model_opts) else: print(f"[I] Load SD3ControlNetModel model from: {self.controlnet_model_dir}") model = SD3ControlNetModel.from_pretrained(self.controlnet_model_dir, **model_opts).to(self.device) # Load transformer model for pos_embed transformer = SD3Transformer2DModel.from_pretrained(self.transformer_model_dir, **model_opts).to(self.device) if hasattr(model.config, "use_pos_embed") and model.config.use_pos_embed is False: pos_embed = model._get_pos_embed_from_transformer(transformer) model.pos_embed = pos_embed.to(model.dtype).to(model.device) # Free transformer model del transformer model = optimizer.optimize_checkpoint(model, torch_inference) model = SD3ControlNetWrapper(model) return model def get_input_names(self): return ["hidden_states", "controlnet_cond", "conditioning_scale", "pooled_projections", "timestep"] def get_output_names(self): return ["controlnet_block_samples"] def get_dynamic_axes(self): xB = "2B" if self.xB == 2 else "B" dynamic_axes = { "hidden_states": {0: xB, 2: "H", 3: "W"}, "controlnet_cond": {0: xB, 2: "H", 3: "W"}, "pooled_projections": {0: xB}, "timestep": {0: xB}, } return dynamic_axes def get_input_profile( self, batch_size: int, image_height: int, image_width: int, static_batch: bool, static_shape: bool, ): latent_height, latent_width = self.check_dims(batch_size, image_height, image_width) ( min_batch, max_batch, _, _, _, _, min_latent_height, max_latent_height, min_latent_width, max_latent_width, ) = self.get_minmax_dims(batch_size, image_height, image_width, static_batch, static_shape) input_profile = { "hidden_states": [ (self.xB * min_batch, self.config["in_channels"], min_latent_height, min_latent_width), (self.xB * batch_size, self.config["in_channels"], latent_height, latent_width), (self.xB * max_batch, self.config["in_channels"], max_latent_height, max_latent_width), ], "timestep": [(self.xB * min_batch,), (self.xB * batch_size,), (self.xB * max_batch,)], "pooled_projections": [ (self.xB * min_batch, self.config["pooled_projection_dim"]), (self.xB * batch_size, self.config["pooled_projection_dim"]), (self.xB * max_batch, self.config["pooled_projection_dim"]), ], "controlnet_cond": [ (self.xB * min_batch, self.config["in_channels"], min_latent_height, min_latent_width), (self.xB * batch_size, self.config["in_channels"], latent_height, latent_width), (self.xB * max_batch, self.config["in_channels"], max_latent_height, max_latent_width), ], } return input_profile def get_shape_dict(self, batch_size, image_height, image_width): latent_height, latent_width = self.check_dims(batch_size, image_height, image_width) shape_dict = { "hidden_states": (self.xB * batch_size, self.config["in_channels"], latent_height, latent_width), "timestep": (self.xB * batch_size,), "pooled_projections": (self.xB * batch_size, self.config["pooled_projection_dim"]), "controlnet_cond": (self.xB * batch_size, self.config["in_channels"], latent_height, latent_width), "conditioning_scale": (), "controlnet_block_samples": ( self.config["num_layers"], self.xB * batch_size, latent_height // 2 * latent_width // 2, self.config["num_attention_heads"] * self.config["attention_head_dim"], ), } return shape_dict def get_sample_input(self, batch_size, image_height, image_width, static_shape): dtype = torch.float16 if self.fp16 else torch.bfloat16 if self.bf16 else torch.float32 latent_height, latent_width = self.check_dims(batch_size, image_height, image_width) sample_input = ( torch.randn( self.xB * batch_size, self.config["in_channels"], latent_height, latent_width, dtype=dtype, device=self.device, ), torch.randn( self.xB * batch_size, self.config["in_channels"], latent_height, latent_width, dtype=dtype, device=self.device, ), torch.tensor(1.0, dtype=dtype, device=self.device), torch.randn(self.xB * batch_size, self.config["pooled_projection_dim"], dtype=dtype, device=self.device), torch.randn(self.xB * batch_size, dtype=torch.float32, device=self.device), ) return sample_input