import os os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") import spaces # noqa: F401 must precede torch / diffusers import json import tempfile from pathlib import Path import gradio as gr import torch from torch.utils._python_dispatch import is_traceable_wrapper_subclass, transform_subclass # --------------------------------------------------------------------------- # ZeroGPU + NVFP4 interop patches # # Cosmos3-Super is 64B params -> ~131 GB at BF16, which does not fit ZeroGPU's # 96 GB xlarge slice. NVFP4 weight-only quantization brings the transformer to # ~36 GB, but two things break along the way and both need patching before the # model is loaded. # --------------------------------------------------------------------------- # (1) ZeroGPU's empty_fake calls empty_like + set_ on each parameter to build the # pinned-CPU mirror it streams from. Those ops don't make sense on tensor-subclass # wrappers (NVFP4Tensor, etc.) which contain multiple inner storages. Patch # empty_fake to recurse into wrapper subclasses via transform_subclass so each # inner tensor gets packed individually. import spaces.zero.torch.patching as _zg_patching _orig_empty_fake = _zg_patching.empty_fake def _empty_fake_subclass_aware(tensor): if is_traceable_wrapper_subclass(tensor): def _per_inner(_name, inner): inner_fake = _orig_empty_fake(inner) # Register inner-tensor aliases so the packer actually packs each storage. _zg_patching.cuda_aliases[inner_fake] = inner return inner_fake return transform_subclass(tensor, _per_inner) return _orig_empty_fake(tensor) _zg_patching.empty_fake = _empty_fake_subclass_aware from diffusers import AutoModel, Cosmos3OmniPipeline, TorchAoConfig from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler from torchao.prototype.mx_formats import NVFP4WeightOnlyConfig from torchao.prototype.mx_formats.nvfp4_tensor import NVFP4Tensor # (2) Cosmos3's time_proj emits fp32 sinusoidals; vanilla F.linear upcasts the # weight, but the NVFP4 dispatch handlers expect input.dtype == weight.orig_dtype. # Wrap the matmul-family handlers to cast non-NVFP4 tensor inputs to the weight's # orig_dtype on the fly. def _make_dtype_safe(orig_handler): def wrapped(func, types, args, kwargs): weight = next((a for a in args if isinstance(a, NVFP4Tensor)), None) if weight is not None: target = weight.orig_dtype new_args = tuple( a.to(target) if isinstance(a, torch.Tensor) and not isinstance(a, NVFP4Tensor) and a.dtype != target and a.is_floating_point() else a for a in args ) return orig_handler(func, types, new_args, kwargs) return orig_handler(func, types, args, kwargs) return wrapped _aten = torch.ops.aten _nvfp4_table = NVFP4Tensor._ATEN_OP_TABLE[NVFP4Tensor] for _f in [ torch.nn.functional.linear, _aten.linear.default, _aten.addmm.default, _aten.mm.default, _aten.matmul.default, ]: if _f in _nvfp4_table: _nvfp4_table[_f] = _make_dtype_safe(_nvfp4_table[_f]) # --------------------------------------------------------------------------- # Model # --------------------------------------------------------------------------- MODEL_ID = "nvidia/Cosmos3-Super-Text2Image" transformer = AutoModel.from_pretrained( MODEL_ID, subfolder="transformer", quantization_config=TorchAoConfig(NVFP4WeightOnlyConfig()), torch_dtype=torch.bfloat16, ) pipe = Cosmos3OmniPipeline.from_pretrained( MODEL_ID, transformer=transformer, torch_dtype=torch.bfloat16, enable_safety_checker=False, ) # Keep the pristine scheduler config around so flow_shift can be re-applied per call. BASE_SCHEDULER_CONFIG = dict(pipe.scheduler.config) pipe.scheduler = UniPCMultistepScheduler.from_config(BASE_SCHEDULER_CONFIG, flow_shift=3.0) pipe.to("cuda") # Resolutions the model card lists as supported: 256p / 480p / 720p at # 16:9, 4:3, 1:1, 3:4, 9:16. 1024x1024 is what NVIDIA's own diffusers snippet uses. RESOLUTIONS = { "1024 x 1024 (1:1)": (1024, 1024), "1280 x 720 (16:9, 720p)": (1280, 720), "720 x 1280 (9:16, 720p)": (720, 1280), "960 x 720 (4:3, 720p)": (960, 720), "720 x 960 (3:4, 720p)": (720, 960), "720 x 720 (1:1, 720p)": (720, 720), "854 x 480 (16:9, 480p)": (854, 480), "480 x 854 (9:16, 480p)": (480, 854), "480 x 480 (1:1, 480p)": (480, 480), "456 x 256 (16:9, 256p)": (456, 256), "256 x 256 (1:1, 256p)": (256, 256), } DEFAULT_RESOLUTION = "1024 x 1024 (1:1)" # NVIDIA's own example prompt for this checkpoint (assets/original_prompt.txt), # plus Physical AI prompts in the robotics / AV / industrial domains the model # was built for. EXAMPLE_PROMPTS = [ "Photorealistic studio photograph of a pottery wheel in motion, a cylinder of wet gray clay " "spinning with concentric rings. Two hands, damp and coated with slip, gently pinch and pull " "the walls upward to form a narrow neck and rounded belly like a vase. Directional warm studio " "light from the right highlights the sheen of water on clay and the texture of fingerprints. " "Splattered clay dots the black apron and wheel tray. Camera slightly above rim height, 35mm, " "crisp focus on hands and clay, background tools softly blurred.", "A warehouse robot with a two-finger gripper folds a blue cloth on a stainless steel workbench, " "overhead fluorescent lighting, shallow depth of field, photorealistic industrial documentation.", "Front-facing camera view from an autonomous vehicle on a rain-slicked city street at dusk, " "wet asphalt reflecting red tail lights, pedestrians crossing at a marked crosswalk, " "photorealistic, 28mm lens.", "A Franka Panda robotic arm picking a ripe tomato from a vertical hydroponic rack in a research " "greenhouse, soft diffused daylight through polycarbonate panels, crisp focus on the gripper.", ] def _parse_prompt(prompt: str) -> str: """Pass JSON-upsampled prompts through as compact JSON, plain text unchanged.""" stripped = prompt.strip() if stripped.startswith("{") and stripped.endswith("}"): try: return json.dumps(json.loads(stripped)) except json.JSONDecodeError: return stripped return stripped def _duration(prompt, resolution=DEFAULT_RESOLUTION, steps=35, *_): width, height = RESOLUTIONS[resolution] # Measured: ~14s per step at 1024x1024 NVFP4 dequant; scale by pixel count + 25% margin. per_step = 18 * (width * height) / (1024 * 1024) return min(1500, int(60 + per_step * int(steps))) @spaces.GPU(duration=_duration, size="xlarge") def generate( prompt: str, resolution: str = DEFAULT_RESOLUTION, steps: int = 35, guidance: float = 4.0, flow_shift: float = 3.0, negative_prompt: str = "", seed: int = 0, randomize_seed: bool = True, progress=gr.Progress(track_tqdm=True), ): """Generate an image from a text prompt with NVIDIA Cosmos3-Super-Text2Image (64B). Args: prompt: Scene description. Plain text, or a JSON-upsampled prompt in NVIDIA's structured format (subjects / lighting / cinematography / ...) for best quality. resolution: Output size and aspect ratio, e.g. "1280 x 720 (16:9, 720p)". steps: Denoising steps. NVIDIA recommends 35-50. guidance: Classifier-free guidance scale. NVIDIA recommends 4.0-6.0. flow_shift: Flow-matching shift for the UniPC scheduler. NVIDIA recommends 3.0-10.0. negative_prompt: Attributes to steer away from. Leave empty to disable. seed: Random seed for reproducibility. randomize_seed: Draw a fresh random seed instead of using `seed`. Returns: The generated PNG image, and the seed that produced it. """ if not prompt or not prompt.strip(): raise gr.Error("Please enter a prompt.") width, height = RESOLUTIONS[resolution] if randomize_seed: seed = int(torch.randint(0, 2**31 - 1, (1,)).item()) generator = torch.Generator(device="cuda").manual_seed(int(seed)) pipe.scheduler = UniPCMultistepScheduler.from_config( BASE_SCHEDULER_CONFIG, flow_shift=float(flow_shift) ) result = pipe( prompt=_parse_prompt(prompt), negative_prompt=negative_prompt or None, num_frames=1, height=height, width=width, num_inference_steps=int(steps), guidance_scale=float(guidance), generator=generator, output_type="pil", ) image = result.video[0] out_path = Path(tempfile.mkdtemp(prefix="cosmos3_")) / "image.png" image.save(out_path) return str(out_path), seed CSS = """ .gradio-container { max-width: 1100px !important; margin: auto !important; } """ with gr.Blocks(theme=gr.themes.Soft(), css=CSS, title="Cosmos3-Super-Text2Image") as demo: gr.Markdown( "# 🌌 Cosmos3-Super-Text2Image\n" "[**nvidia/Cosmos3-Super-Text2Image**](https://huggingface.co/nvidia/Cosmos3-Super-Text2Image) " "— a 64B-parameter omnimodal world model for Physical AI, generating high-fidelity images " "from text. Running on a single Blackwell GPU via NVFP4 weight-only quantization.\n\n" "*NVIDIA officially tests this checkpoint only at BF16 (4xH200 / 8xH100). NVFP4 is " "unofficial and may show quality drift versus the full-precision recipe.*" ) with gr.Row(): prompt = gr.Textbox( show_label=False, placeholder="A warehouse robot folds a blue cloth on a clean workbench...", container=False, scale=4, lines=2, ) run = gr.Button("Generate", variant="primary", scale=1) out = gr.Image(label="Output", type="filepath", format="png", height=640) with gr.Accordion("Advanced settings", open=False): negative_prompt = gr.Textbox(label="Negative prompt", value="") resolution = gr.Dropdown( label="Resolution", choices=list(RESOLUTIONS), value=DEFAULT_RESOLUTION, ) with gr.Row(): steps = gr.Slider(label="Inference steps", minimum=10, maximum=50, value=35, step=1) guidance = gr.Slider( label="Guidance scale", minimum=1.0, maximum=8.0, value=4.0, step=0.1 ) flow_shift = gr.Slider( label="Flow shift", minimum=1.0, maximum=12.0, value=3.0, step=0.5 ) with gr.Row(): randomize_seed = gr.Checkbox(label="Randomize seed", value=True) seed = gr.Number(label="Seed", value=0, precision=0) gr.Markdown( "**Tip:** NVIDIA recommends upsampling prompts into a structured JSON format " "(subjects, lighting, cinematography, quadrant scan) for best quality — see the " "[prompt upsampling guide](https://github.com/nvidia/cosmos-framework/blob/main/docs/prompt_upsampling.md). " "Paste that JSON straight into the prompt box and it is forwarded as-is." ) gr.Examples( examples=[[p] for p in EXAMPLE_PROMPTS], inputs=[prompt], outputs=[out, seed], fn=generate, cache_examples=True, cache_mode="lazy", ) inputs = [ prompt, resolution, steps, guidance, flow_shift, negative_prompt, seed, randomize_seed, ] outputs = [out, seed] run.click(generate, inputs, outputs) prompt.submit(generate, inputs, outputs) demo.queue().launch(mcp_server=True)