1import json
2import torch
3from diffusers import Cosmos3OmniPipeline, Cosmos3OmniTransformer
4from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler
5from modelopt.torch.quantization.qtensor.base_qtensor import QTensorWrapper
6import modelopt.torch.opt as mto
7
8
9def patch_modelopt_qtensor_loader():
10 import accelerate.utils.modeling as accelerate_modeling
11 import diffusers.models.model_loading_utils as diffusers_loading
12
13 original = accelerate_modeling.set_module_tensor_to_device
14 if getattr(original, "_cosmos3_modelopt_patch", False):
15 return
16
17 def patched(module, tensor_name, device, value=None, dtype=None, fp16_statistics=None,
18 tied_params_map=None, non_blocking=False, clear_cache=True):
19 leaf_module = module
20 leaf_name = tensor_name
21 if "." in tensor_name:
22 parts = tensor_name.split(".")
23 for part in parts[:-1]:
24 leaf_module = getattr(leaf_module, part)
25 leaf_name = parts[-1]
26 old_value = getattr(leaf_module, leaf_name) if hasattr(leaf_module, leaf_name) else None
27 if isinstance(old_value, QTensorWrapper) and value is not None:
28 leaf_module._parameters[leaf_name] = QTensorWrapper(
29 value.to(device, non_blocking=non_blocking),
30 metadata=old_value.metadata,
31 )
32 return
33 return original(module, tensor_name, device, value, dtype, fp16_statistics,
34 tied_params_map, non_blocking, clear_cache)
35
36 patched._cosmos3_modelopt_patch = True
37 accelerate_modeling.set_module_tensor_to_device = patched
38 diffusers_loading.set_module_tensor_to_device = patched
39
40
41def cast_modelopt_runtime_tensors(model, dtype=torch.bfloat16):
42 for module in model.modules():
43 for name, param in list(module._parameters.items()):
44 if isinstance(param, QTensorWrapper):
45 param.metadata["dtype"] = dtype
46 elif param is not None and param.is_floating_point():
47 module._parameters[name] = torch.nn.Parameter(
48 param.detach().to(dtype),
49 requires_grad=param.requires_grad,
50 )
51 for name, buf in list(module._buffers.items()):
52 if buf is not None and buf.is_floating_point():
53 module._buffers[name] = buf.to(dtype)
54 return model
55
56
57patch_modelopt_qtensor_loader()
58mto.enable_huggingface_checkpointing()
59
60transformer = Cosmos3OmniTransformer.from_pretrained(
61 "WaveCut/Cosmos3-Super-Text2Image-ModelOpt-FP8-Transformer",
62 subfolder="transformer",
63 use_safetensors=False,
64)
65transformer = cast_modelopt_runtime_tensors(transformer, torch.bfloat16)
66
67pipe = Cosmos3OmniPipeline.from_pretrained(
68 "nvidia/Cosmos3-Super-Text2Image",
69 transformer=transformer,
70 torch_dtype=torch.bfloat16,
71 device_map="cuda",
72 enable_safety_checker=True,
73)
74pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=3.0)
75pipe.to("cuda")
76
77json_caption = {
78 "subjects": [],
79 "background_setting": "A concise scene description.",
80 "comprehensive_t2i_caption": "A detailed natural-language caption.",
81 "resolution": {"H": 1024, "W": 1024},
82 "aspect_ratio": "1,1",
83}
84
85with torch.autocast("cuda", dtype=torch.bfloat16):
86 result = pipe(
87 prompt=json.dumps(json_caption),
88 negative_prompt="",
89 num_frames=1,
90 height=1024,
91 width=1024,
92 num_inference_steps=50,
93 guidance_scale=4.0,
94 generator=torch.Generator(device="cuda").manual_seed(1143),
95 )
96result.video[0].save("cosmos3_modelopt_fp8.png")