# # 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 huggingface_hub import hf_hub_download from safetensors import safe_open from demo_diffusion.dynamic_import import import_from_diffusers from demo_diffusion.model import base_model, load, optimizer from demo_diffusion.utils_sd3.other_impls import load_into from demo_diffusion.utils_sd3.sd3_impls import BaseModel as BaseModelSD3 # List of models to import from diffusers.models models_to_import = ["FluxTransformer2DModel", "SD3Transformer2DModel", "WanTransformer3DModel", "CosmosTransformer3DModel"] for model in models_to_import: globals()[model] = import_from_diffusers(model, "diffusers.models") # Import FluxKontextUtil from pipeline module # Using a deferred import to avoid circular dependencies def _get_flux_kontext_util(): from demo_diffusion.pipeline.flux_pipeline import FluxKontextUtil return FluxKontextUtil class SD3_MMDiTModel(base_model.BaseModel): def __init__( self, version, pipeline, device, hf_token, verbose, framework_model_dir, shift=1.0, fp16=False, max_batch_size=16, text_maxlen=77, ): super(SD3_MMDiTModel, self).__init__( version, pipeline, device=device, hf_token=hf_token, verbose=verbose, framework_model_dir=framework_model_dir, fp16=fp16, max_batch_size=max_batch_size, text_maxlen=text_maxlen, ) self.subfolder = "sd3" self.mmdit_dim = 16 self.shift = shift self.xB = 2 def get_model(self, torch_inference=""): sd3_model_dir = load.get_checkpoint_dir(self.framework_model_dir, self.version, self.pipeline, self.subfolder) sd3_filename = "sd3_medium.safetensors" sd3_model_path = f"{sd3_model_dir}/{sd3_filename}" if not os.path.exists(sd3_model_path): hf_hub_download(repo_id=self.path, filename=sd3_filename, local_dir=sd3_model_dir) with safe_open(sd3_model_path, framework="pt", device=self.device) as f: model = BaseModelSD3( shift=self.shift, file=f, prefix="model.diffusion_model.", device=self.device, dtype=torch.float16 ).eval() load_into(f, model, "model.", self.device, torch.float16) model = optimizer.optimize_checkpoint(model, torch_inference) return model def get_input_names(self): return ["sample", "sigma", "c_crossattn", "y"] def get_output_names(self): return ["latent"] def get_dynamic_axes(self): xB = "2B" if self.xB == 2 else "B" return { "sample": {0: xB, 2: "H", 3: "W"}, "sigma": {0: xB}, "c_crossattn": {0: xB}, "y": {0: xB}, "latent": {0: xB, 2: "H", 3: "W"}, } def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape): 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) ) return { "sample": [ (self.xB * min_batch, self.mmdit_dim, min_latent_height, min_latent_width), (self.xB * batch_size, self.mmdit_dim, latent_height, latent_width), (self.xB * max_batch, self.mmdit_dim, max_latent_height, max_latent_width), ], "sigma": [(self.xB * min_batch,), (self.xB * batch_size,), (self.xB * max_batch,)], "c_crossattn": [ (self.xB * min_batch, 154, 4096), (self.xB * batch_size, 154, 4096), (self.xB * max_batch, 154, 4096), ], "y": [(self.xB * min_batch, 2048), (self.xB * batch_size, 2048), (self.xB * max_batch, 2048)], } def get_shape_dict(self, batch_size, image_height, image_width): latent_height, latent_width = self.check_dims(batch_size, image_height, image_width) return { "sample": (self.xB * batch_size, self.mmdit_dim, latent_height, latent_width), "sigma": (self.xB * batch_size,), "c_crossattn": (self.xB * batch_size, 154, 4096), "y": (self.xB * batch_size, 2048), "latent": (self.xB * batch_size, self.mmdit_dim, latent_height, latent_width), } def get_sample_input(self, batch_size, image_height, image_width, static_shape): latent_height, latent_width = self.check_dims(batch_size, image_height, image_width) dtype = torch.float16 if self.fp16 else torch.float32 return ( torch.randn(batch_size, self.mmdit_dim, latent_height, latent_width, dtype=dtype, device=self.device), torch.randn(batch_size, dtype=dtype, device=self.device), { "c_crossattn": torch.randn(batch_size, 154, 4096, dtype=dtype, device=self.device), "y": torch.randn(batch_size, 2048, dtype=dtype, device=self.device), }, ) class FluxTransformerModel(base_model.BaseModel): def __init__( self, version, pipeline, device, hf_token, verbose, framework_model_dir, fp16=False, tf32=False, int8=False, fp8=False, bf16=False, max_batch_size=16, text_maxlen=77, weight_streaming=False, weight_streaming_budget_percentage=None, kontext_resolution=None, ): super(FluxTransformerModel, self).__init__( version, pipeline, device=device, hf_token=hf_token, verbose=verbose, framework_model_dir=framework_model_dir, fp16=fp16, tf32=tf32, int8=int8, fp8=fp8, bf16=bf16, max_batch_size=max_batch_size, text_maxlen=text_maxlen, ) self.subfolder = "transformer" self.transformer_model_dir = load.get_checkpoint_dir( self.framework_model_dir, self.version, self.pipeline, self.subfolder ) if not os.path.exists(self.transformer_model_dir): self.config = FluxTransformer2DModel.load_config(self.path, subfolder=self.subfolder, token=self.hf_token) else: print(f"[I] Load FluxTransformer2DModel config from: {self.transformer_model_dir}") self.config = FluxTransformer2DModel.load_config(self.transformer_model_dir) self.weight_streaming = weight_streaming self.weight_streaming_budget_percentage = weight_streaming_budget_percentage self.out_channels = self.config.get("out_channels") or self.config["in_channels"] self.kontext_resolution = kontext_resolution 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.transformer_model_dir, model_opts, self.hf_safetensor): model = FluxTransformer2DModel.from_pretrained( self.path, subfolder=self.subfolder, use_safetensors=self.hf_safetensor, token=self.hf_token, **model_opts, ).to(self.device) model.save_pretrained(self.transformer_model_dir, **model_opts) else: print(f"[I] Load FluxTransformer2DModel model from: {self.transformer_model_dir}") model = FluxTransformer2DModel.from_pretrained(self.transformer_model_dir, **model_opts).to(self.device) if torch_inference: model.to(memory_format=torch.channels_last) model = optimizer.optimize_checkpoint(model, torch_inference) return model def get_input_names(self): return [ "hidden_states", "encoder_hidden_states", "pooled_projections", "timestep", "img_ids", "txt_ids", "guidance", ] def get_output_names(self): return ["latent"] def get_dynamic_axes(self): dynamic_axes = { "hidden_states": {0: "B", 1: "latent_dim"}, "encoder_hidden_states": {0: "B"}, "pooled_projections": {0: "B"}, "timestep": {0: "B"}, "img_ids": {0: "latent_dim"}, "txt_ids": {}, } if self.config["guidance_embeds"]: dynamic_axes["guidance"] = {0: "B"} return dynamic_axes def get_context_latent_dim(self, static_shape=False): FluxKontextUtil = _get_flux_kontext_util() return FluxKontextUtil.get_context_latent_dim( version=self.version, kontext_resolution=self.kontext_resolution, compression_factor=self.compression_factor, static_shape=static_shape, ) def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape): latent_height, latent_width = self.check_dims(batch_size, image_height, image_width) ( min_batch, max_batch, min_image_height, max_image_height, min_image_width, max_image_width, 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) min_context_latent_dim, context_latent_dim, max_context_latent_dim = self.get_context_latent_dim(static_shape) input_profile = { "hidden_states": [ ( min_batch, (min_latent_height // 2) * (min_latent_width // 2) + min_context_latent_dim, self.config["in_channels"], ), ( batch_size, (latent_height // 2) * (latent_width // 2) + context_latent_dim, self.config["in_channels"], ), ( max_batch, (max_latent_height // 2) * (max_latent_width // 2) + max_context_latent_dim, self.config["in_channels"], ), ], "encoder_hidden_states": [ (min_batch, self.text_maxlen, self.config["joint_attention_dim"]), (batch_size, self.text_maxlen, self.config["joint_attention_dim"]), (max_batch, self.text_maxlen, self.config["joint_attention_dim"]), ], "pooled_projections": [ (min_batch, self.config["pooled_projection_dim"]), (batch_size, self.config["pooled_projection_dim"]), (max_batch, self.config["pooled_projection_dim"]), ], "timestep": [(min_batch,), (batch_size,), (max_batch,)], "img_ids": [ ((min_latent_height // 2) * (min_latent_width // 2) + min_context_latent_dim, 3), ((latent_height // 2) * (latent_width // 2) + context_latent_dim, 3), ((max_latent_height // 2) * (max_latent_width // 2) + max_context_latent_dim, 3), ], "txt_ids": [(self.text_maxlen, 3), (self.text_maxlen, 3), (self.text_maxlen, 3)], } if self.config["guidance_embeds"]: input_profile["guidance"] = [(min_batch,), (batch_size,), (max_batch,)] 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) _, context_latent_dim, _ = self.get_context_latent_dim() shape_dict = { "hidden_states": ( batch_size, (latent_height // 2) * (latent_width // 2) + context_latent_dim, self.config["in_channels"], ), "encoder_hidden_states": (batch_size, self.text_maxlen, self.config["joint_attention_dim"]), "pooled_projections": (batch_size, self.config["pooled_projection_dim"]), "timestep": (batch_size,), "img_ids": ((latent_height // 2) * (latent_width // 2) + context_latent_dim, 3), "txt_ids": (self.text_maxlen, 3), "latent": (batch_size, (latent_height // 2) * (latent_width // 2) + context_latent_dim, self.out_channels), } if self.config["guidance_embeds"]: shape_dict["guidance"] = (batch_size,) return shape_dict def get_sample_input(self, batch_size, image_height, image_width, static_shape): latent_height, latent_width = self.check_dims(batch_size, image_height, image_width) dtype = torch.float32 assert not (self.fp16 and self.bf16), "fp16 and bf16 cannot be enabled simultaneously" tensor_dtype = torch.bfloat16 if self.bf16 else (torch.float16 if self.fp16 else torch.float32) sample_input = ( torch.randn( batch_size, (latent_height // 2) * (latent_width // 2), self.config["in_channels"], dtype=tensor_dtype, device=self.device, ), torch.randn( batch_size, self.text_maxlen, self.config["joint_attention_dim"], dtype=tensor_dtype, device=self.device ), torch.randn(batch_size, self.config["pooled_projection_dim"], dtype=tensor_dtype, device=self.device), torch.tensor([1.0] * batch_size, dtype=tensor_dtype, device=self.device), torch.randn((latent_height // 2) * (latent_width // 2), 3, dtype=dtype, device=self.device), torch.randn(self.text_maxlen, 3, dtype=dtype, device=self.device), {}, ) if self.config["guidance_embeds"]: sample_input[-1]["guidance"] = torch.tensor([1.0] * batch_size, dtype=dtype, device=self.device) return sample_input def optimize(self, onnx_graph): if self.fp8: return super().optimize(onnx_graph) if self.int8: return super().optimize(onnx_graph, modify_int8_graph=True) return super().optimize(onnx_graph) class UpcastLayer(torch.nn.Module): def __init__(self, base_layer: torch.nn.Module, upcast_to: torch.dtype): super().__init__() self.output_dtype = next(base_layer.parameters()).dtype self.upcast_to = upcast_to self.context_pre_only = base_layer.context_pre_only base_layer = base_layer.to(dtype=self.upcast_to) self.base_layer = base_layer def forward(self, *inputs, **kwargs): casted_inputs = tuple( in_val.to(self.upcast_to) if isinstance(in_val, torch.Tensor) else in_val for in_val in inputs ) kwarg_casted = {} for name, val in kwargs.items(): kwarg_casted[name] = val.to(dtype=self.upcast_to) if isinstance(val, torch.Tensor) else val output = self.base_layer(*casted_inputs, **kwarg_casted) if isinstance(output, tuple): output = tuple(out.to(self.output_dtype) if isinstance(out, torch.Tensor) else out for out in output) else: output = output.to(dtype=self.output_dtype) return output class SD3TransformerModel(base_model.BaseModel): def __init__( self, version, pipeline, device, hf_token, verbose, framework_model_dir, fp16=False, tf32=False, bf16=False, fp8=False, int8=False, fp4=False, max_batch_size=16, text_maxlen=256, weight_streaming=False, weight_streaming_budget_percentage=None, do_classifier_free_guidance=False, ): super(SD3TransformerModel, self).__init__( version, pipeline, device=device, hf_token=hf_token, verbose=verbose, framework_model_dir=framework_model_dir, fp16=fp16, tf32=tf32, bf16=bf16, fp8=fp8, int8=int8, fp4=fp4, max_batch_size=max_batch_size, text_maxlen=text_maxlen, ) self.subfolder = "transformer" self.transformer_model_dir = load.get_checkpoint_dir( self.framework_model_dir, self.version, self.pipeline, self.subfolder ) if not os.path.exists(self.transformer_model_dir): self.config = SD3Transformer2DModel.load_config(self.path, subfolder=self.subfolder, token=self.hf_token) else: print(f"[I] Load SD3Transformer2DModel config from: {self.transformer_model_dir}") self.config = SD3Transformer2DModel.load_config(self.transformer_model_dir) self.weight_streaming = weight_streaming self.weight_streaming_budget_percentage = weight_streaming_budget_percentage self.out_channels = self.config.get("out_channels") self.xB = 2 if do_classifier_free_guidance else 1 # batch multiplier self.num_controlnet_layers = 19 # Can be queried from the ControlNet model config 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.transformer_model_dir, model_opts, self.hf_safetensor): model = SD3Transformer2DModel.from_pretrained( self.path, subfolder=self.subfolder, use_safetensors=self.hf_safetensor, token=self.hf_token, **model_opts, ).to(self.device) model.save_pretrained(self.transformer_model_dir, **model_opts) else: print(f"[I] Load SD3Transformer2DModel model from: {self.transformer_model_dir}") model = SD3Transformer2DModel.from_pretrained(self.transformer_model_dir, **model_opts).to(self.device) if self.version == "3.5-large": model.transformer_blocks[35] = UpcastLayer(model.transformer_blocks[35], torch.float32) if torch_inference: model.to(memory_format=torch.channels_last) model = optimizer.optimize_checkpoint(model, torch_inference) return model def get_input_names(self): input_names = [ "hidden_states", "encoder_hidden_states", "pooled_projections", "timestep", "block_controlnet_hidden_states" ] return input_names def get_output_names(self): return ["latent"] def get_dynamic_axes(self): xB = "2B" if self.xB == 2 else "B" dynamic_axes = { "hidden_states": {0: xB, 2: "H", 3: "W"}, "encoder_hidden_states": {0: xB}, "pooled_projections": {0: xB}, "timestep": {0: xB}, "latent": {0: xB, 2: "H", 3: "W"}, "block_controlnet_hidden_states": {1: xB, 2: "latent_dim"} } 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), ], "encoder_hidden_states": [ (self.xB * min_batch, self.text_maxlen, self.config["joint_attention_dim"]), (self.xB * batch_size, self.text_maxlen, self.config["joint_attention_dim"]), (self.xB * max_batch, self.text_maxlen, self.config["joint_attention_dim"]), ], "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"]), ], "timestep": [(self.xB * min_batch,), (self.xB * batch_size,), (self.xB * max_batch,)], "block_controlnet_hidden_states": [ ( self.num_controlnet_layers, self.xB * min_batch, min_latent_height // self.config["patch_size"] * min_latent_width // self.config["patch_size"], self.config["num_attention_heads"] * self.config["attention_head_dim"], ), ( self.num_controlnet_layers, self.xB * batch_size, latent_height // self.config["patch_size"] * latent_width // self.config["patch_size"], self.config["num_attention_heads"] * self.config["attention_head_dim"], ), ( self.num_controlnet_layers, self.xB * max_batch, max_latent_height // self.config["patch_size"] * max_latent_width // self.config["patch_size"], self.config["num_attention_heads"] * self.config["attention_head_dim"], ), ] } 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), "encoder_hidden_states": (self.xB * batch_size, self.text_maxlen, self.config["joint_attention_dim"]), "pooled_projections": (self.xB * batch_size, self.config["pooled_projection_dim"]), "timestep": (self.xB * batch_size,), "latent": (self.xB * batch_size, self.out_channels, latent_height, latent_width), "block_controlnet_hidden_states": ( self.num_controlnet_layers, self.xB * batch_size, latent_height // self.config["patch_size"] * latent_width // self.config["patch_size"], 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): assert not (self.fp16 and self.bf16), "fp16 and bf16 cannot be enabled simultaneously" 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.text_maxlen, self.config["joint_attention_dim"], 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), { "block_controlnet_hidden_states": torch.randn( self.num_controlnet_layers, self.xB * batch_size, latent_height // self.config["patch_size"] * latent_width // self.config["patch_size"], self.config["num_attention_heads"] * self.config["attention_head_dim"], dtype=dtype, device=self.device, ), } ) return sample_input class WanTransformerModel(base_model.BaseModel): def __init__( self, version, pipeline, device, hf_token, verbose, framework_model_dir, subfolder="transformer", fp16=False, tf32=False, bf16=True, fp8=False, int8=False, max_batch_size=1, text_maxlen=512, num_frames=81, height=720, width=1280, weight_streaming=False, weight_streaming_budget_percentage=None, ): super(WanTransformerModel, self).__init__( version, pipeline, device=device, hf_token=hf_token, verbose=verbose, framework_model_dir=framework_model_dir, fp16=fp16, tf32=tf32, bf16=bf16, fp8=fp8, int8=int8, max_batch_size=max_batch_size, text_maxlen=text_maxlen, embedding_dim=4096, compression_factor=8, ) self.subfolder = subfolder self.transformer_model_dir = load.get_checkpoint_dir( self.framework_model_dir, self.version, self.pipeline, self.subfolder ) if not os.path.exists(self.transformer_model_dir): self.config = WanTransformer3DModel.load_config( self.path, subfolder=self.subfolder, token=self.hf_token ) else: print(f"[I] Load WanTransformer3DModel config from: {self.transformer_model_dir}") self.config = WanTransformer3DModel.load_config(self.transformer_model_dir) self.weight_streaming = weight_streaming self.weight_streaming_budget_percentage = weight_streaming_budget_percentage self.do_constant_folding = False self.latent_channels = self.config.get("in_channels", 16) self.temporal_compression_factor = 4 self.num_frames = num_frames self.min_latent_frames = 81 # hardcode to 81 frames for Wan 2.2 self.max_latent_frames = 81 def get_model(self, torch_inference=""): model_opts = {"torch_dtype": torch.bfloat16} if self.bf16 else {} if not load.is_model_cached(self.transformer_model_dir, model_opts, self.hf_safetensor): model = WanTransformer3DModel.from_pretrained( self.path, subfolder=self.subfolder, use_safetensors=self.hf_safetensor, token=self.hf_token, **model_opts, ).to(self.device) model.save_pretrained(self.transformer_model_dir, **model_opts) else: print(f"[I] Load WanTransformer3DModel model from: {self.transformer_model_dir}") model = WanTransformer3DModel.from_pretrained(self.transformer_model_dir, **model_opts).to(self.device) if torch_inference: model.to(memory_format=torch.channels_last) model = optimizer.optimize_checkpoint(model, torch_inference) return model def get_input_names(self): return [ "hidden_states", "timestep", "encoder_hidden_states", ] def get_output_names(self): return ["denoised_latents"] def get_dynamic_axes(self): return {} def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape, num_frames): latent_height, latent_width, latent_frames = self.check_dims( batch_size, image_height, image_width, num_frames ) ( min_batch, max_batch, _, _, _, _, min_latent_height, max_latent_height, min_latent_width, max_latent_width, min_latent_frames, max_latent_frames ) = self.get_minmax_dims(batch_size, image_height, image_width, static_batch, static_shape, num_frames) input_profile = { "hidden_states": [ (min_batch, self.latent_channels, min_latent_frames, min_latent_height, min_latent_width), (batch_size, self.latent_channels, latent_frames, latent_height, latent_width), (max_batch, self.latent_channels, max_latent_frames, max_latent_height, max_latent_width), ], "timestep": [ (min_batch,), (batch_size,), (max_batch,) ], "encoder_hidden_states": [ (min_batch, self.text_maxlen, self.embedding_dim), (batch_size, self.text_maxlen, self.embedding_dim), (max_batch, self.text_maxlen, self.embedding_dim), ], } return input_profile def get_shape_dict(self, batch_size, image_height, image_width, num_frames): latent_height, latent_width, latent_frames = self.check_dims( batch_size, image_height, image_width, num_frames ) return { "hidden_states": ( batch_size, self.latent_channels, latent_frames, latent_height, latent_width ), "timestep": (batch_size,), "encoder_hidden_states": ( batch_size, self.text_maxlen, self.embedding_dim ), "denoised_latents": ( batch_size, self.latent_channels, latent_frames, latent_height, latent_width ), } def get_sample_input(self, batch_size, image_height, image_width, static_shape, num_frames): latent_height, latent_width, latent_frames = self.check_dims( batch_size, image_height, image_width, num_frames ) assert (self.bf16), "transformer must be BF16" dtype = torch.bfloat16 if self.bf16 else torch.float32 timesteps = torch.tensor([999], dtype=torch.long, device=self.device) if self.subfolder == "transformer" else torch.tensor([1], dtype=torch.long, device=self.device) sample_input = ( torch.randn( batch_size, self.latent_channels, latent_frames, latent_height, latent_width, dtype=dtype, device=self.device, ), timesteps, torch.randn( batch_size, self.text_maxlen, self.embedding_dim, dtype=dtype, device=self.device, ), ) return sample_input class CosmosTransformerModel(base_model.BaseModel): def __init__( self, version, pipeline, device, hf_token, verbose, framework_model_dir, fp16=False, tf32=False, int8=False, fp8=False, bf16=False, max_batch_size=16, text_maxlen=77, weight_streaming=False, weight_streaming_budget_percentage=None, ): super(CosmosTransformerModel, self).__init__( version, pipeline, device=device, hf_token=hf_token, verbose=verbose, framework_model_dir=framework_model_dir, fp16=fp16, tf32=tf32, int8=int8, fp8=fp8, bf16=bf16, max_batch_size=max_batch_size, text_maxlen=text_maxlen, ) self.subfolder = "transformer" self.transformer_model_dir = load.get_checkpoint_dir( self.framework_model_dir, self.version, self.pipeline, self.subfolder ) if not os.path.exists(self.transformer_model_dir): self.config = CosmosTransformer3DModel.load_config(self.path, subfolder=self.subfolder, token=self.hf_token) else: print(f"[I] Load CosmosTransformer3DModel config from: {self.transformer_model_dir}") self.config = CosmosTransformer3DModel.load_config(self.transformer_model_dir) self.weight_streaming = weight_streaming self.weight_streaming_budget_percentage = weight_streaming_budget_percentage 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.transformer_model_dir, model_opts, self.hf_safetensor): model = CosmosTransformer3DModel.from_pretrained( self.path, subfolder=self.subfolder, use_safetensors=self.hf_safetensor, token=self.hf_token, **model_opts, ).to(self.device) model.save_pretrained(self.transformer_model_dir, **model_opts) else: print(f"[I] Load CosmosTransformer3DModel model from: {self.transformer_model_dir}") model = CosmosTransformer3DModel.from_pretrained(self.transformer_model_dir, **model_opts).to(self.device) if torch_inference: model.to(memory_format=torch.channels_last) if self.fp16: model.transformer_blocks[6].attn1.norm_q.float().to(self.device) model = optimizer.optimize_checkpoint(model, torch_inference) return model.to(self.device) def get_input_names(self): input_names = [ "hidden_states", "timestep", "encoder_hidden_states", "padding_mask", ] if self.pipeline_type.is_video2world(): input_names.append("fps") input_names.append("condition_mask") return input_names def get_output_names(self): return ["latent"] def get_dynamic_axes(self): dynamic_axes = { "hidden_states": {0: "B", 2: "latent_frames", 3: "latent_H", 4: "latent_W"}, "timestep": {0: "B"}, "encoder_hidden_states": {0: "B"}, "padding_mask": {0: "B", 2: "H", 3: "W"}, } if self.pipeline_type.is_video2world(): dynamic_axes["fps"] = {0: "B"} dynamic_axes["condition_mask"] = {0: "B", 2: "latent_frames", 3: "latent_H", 4: "latent_W"} return dynamic_axes def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape): latent_height, latent_width = self.check_dims(batch_size, image_height, image_width) ( min_batch, max_batch, min_image_height, max_image_height, min_image_width, max_image_width, 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) latent_frames = 24 if self.pipeline_type.is_video2world() else 1 latent_channels = ( self.config["in_channels"] - 1 if self.pipeline_type.is_video2world() else self.config["in_channels"] ) input_profile = { "hidden_states": [ (min_batch, latent_channels, latent_frames, min_latent_height, min_latent_width), (batch_size, latent_channels, latent_frames, latent_height, latent_width), (max_batch, latent_channels, latent_frames, max_latent_height, max_latent_width), ], "timestep": [(min_batch,), (batch_size,), (max_batch,)], "encoder_hidden_states": [ (min_batch, self.text_maxlen, self.config["text_embed_dim"]), (batch_size, self.text_maxlen, self.config["text_embed_dim"]), (max_batch, self.text_maxlen, self.config["text_embed_dim"]), ], "padding_mask": [ (1, 1, min_image_height, min_image_width), (1, 1, image_height, image_width), (1, 1, max_image_height, max_image_width), ], } if self.pipeline_type.is_video2world(): input_profile["fps"] = [(min_batch,), (batch_size,), (max_batch,)] input_profile["condition_mask"] = [ (min_batch, 1, latent_frames, min_latent_height, min_latent_width), (batch_size, 1, latent_frames, latent_height, latent_width), (max_batch, 1, latent_frames, 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) # TODO: get latent_frames from infer call latent_frames = 24 if self.pipeline_type.is_video2world() else 1 latent_channels = ( self.config["in_channels"] - 1 if self.pipeline_type.is_video2world() else self.config["in_channels"] ) shape_dict = { "hidden_states": (batch_size, latent_channels, latent_frames, latent_height, latent_width), "timestep": (batch_size,), "encoder_hidden_states": (batch_size, self.text_maxlen, self.config["text_embed_dim"]), "padding_mask": (1, 1, image_height, image_width), "latent": (batch_size, self.config["in_channels"], latent_frames, latent_height, latent_width), } if self.pipeline_type.is_video2world(): shape_dict["fps"] = (batch_size,) shape_dict["condition_mask"] = (batch_size, 1, latent_frames, latent_height, latent_width) return shape_dict def get_sample_input(self, batch_size, image_height, image_width, static_shape): latent_height, latent_width = self.check_dims(batch_size, image_height, image_width) dtype = torch.float32 assert not (self.fp16 and self.bf16), "fp16 and bf16 cannot be enabled simultaneously" tensor_dtype = torch.bfloat16 if self.bf16 else (torch.float16 if self.fp16 else torch.float32) latent_frames = 1 latent_channels = ( self.config["in_channels"] - 1 if self.pipeline_type.is_video2world() else self.config["in_channels"] ) sample_input = ( { "hidden_states": torch.randn( batch_size, latent_channels, latent_frames, latent_height, latent_width, dtype=tensor_dtype, device=self.device, ), "timestep": torch.tensor([1.0] * batch_size, dtype=tensor_dtype, device=self.device), "encoder_hidden_states": torch.randn( batch_size, self.text_maxlen, self.config["text_embed_dim"], dtype=tensor_dtype, device=self.device ), "padding_mask": torch.ones( batch_size, 1, image_height, image_width, dtype=tensor_dtype, device=self.device ), }, ) if self.pipeline_type.is_video2world(): sample_input[-1]["fps"] = torch.tensor([30] * batch_size, dtype=dtype, device=self.device) sample_input[-1]["condition_mask"] = torch.randn( batch_size, 1, latent_frames, latent_height, latent_width, dtype=tensor_dtype, device=self.device, ) return sample_input