Upload folder using huggingface_hub

README.md

@@ -1,13 +1,28 @@
 ---
-title: Cosmos3 Super Text2Image
-emoji: 😻
-colorFrom: pink
-colorTo: pink
+title: Cosmos3-Super-Text2Image
+emoji: 🌌
+colorFrom: indigo
+colorTo: purple
 sdk: gradio
-sdk_version: 6.24.0
-python_version: '3.12'
+sdk_version: 6.15.1
 app_file: app.py
+python_version: "3.12"
+short_description: NVIDIA Cosmos3-Super 64B text-to-image, NVFP4 quantization
+startup_duration_timeout: 1h
 pinned: false
+license: other
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# 🌌 Cosmos3-Super-Text2Image
+
+Demo of [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 on a single NVIDIA Blackwell GPU (ZeroGPU) via **NVFP4** weight-only quantization (torchao / NVIDIA ModelOpt).
+
+## Notes
+
+- NVIDIA officially tests this checkpoint **only at BF16**. NVFP4 is unofficial and may show quality drift compared to the full-precision recipe.
+- The 64B transformer does not fit in BF16 on a single GPU, so it is loaded weight-only quantized to NVFP4 and streamed onto the ZeroGPU `xlarge` (96 GB) allocation.
+- The app doubles as an **MCP server** (`mcp_server=True`) — the `generate` tool is exposed with its docstring and type hints.
+
+## License
+
+Model released under [OpenMDW 1.1](https://openmdw.ai/license/1-1/).

app.py

@@ -0,0 +1,222 @@
+import os
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+
+import spaces  # noqa: F401  must precede torch / diffusers
+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 packer support for NVFP4 tensor-subclass weights -----------------
+# 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
+
+# --- NVFP4 dtype-safety shim --------------------------------------------------
+# 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 loading ------------------------------------------------------------
+MODEL_ID = "nvidia/Cosmos3-Super-Text2Image"
+
+quant_config = TorchAoConfig(NVFP4WeightOnlyConfig())
+
+transformer = AutoModel.from_pretrained(
+    MODEL_ID,
+    subfolder="transformer",
+    quantization_config=quant_config,
+    torch_dtype=torch.bfloat16,
+)
+
+pipe = Cosmos3OmniPipeline.from_pretrained(
+    MODEL_ID,
+    transformer=transformer,
+    torch_dtype=torch.bfloat16,
+    enable_safety_checker=False,
+)
+pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=3.0)
+pipe.to("cuda")
+
+RESOLUTIONS = {
+    "1024\u00d71024 (1:1)": (1024, 1024),
+    "1280\u00d7720  (16:9)": (1280, 720),
+    "720\u00d71280  (9:16)": (720, 1280),
+    "1024\u00d7768  (4:3)": (1024, 768),
+    "768\u00d71024  (3:4)": (768, 1024),
+}
+
+
+def _duration(prompt, resolution, steps, *_):
+    w, h = RESOLUTIONS[resolution]
+    # Measured: ~14s per step at 1024\u00d71024 NVFP4 dequant; scale by pixel count + margin.
+    per_step = 18 * (w * h) / (1024 * 1024)
+    return min(1500, int(60 + per_step * int(steps)))
+
+
+@spaces.GPU(duration=_duration, size="xlarge")
+def generate(
+    prompt: str,
+    resolution: str = "1024\u00d71024 (1:1)",
+    steps: int = 35,
+    guidance: float = 4.0,
+    negative_prompt: str = "",
+    seed: int = 0,
+    randomize_seed: bool = True,
+    progress=gr.Progress(track_tqdm=True),
+):
+    """Generate a high-fidelity image from a text prompt with NVIDIA Cosmos3-Super-Text2Image (64B, NVFP4).
+
+    Args:
+        prompt: Text description of the image to generate.
+        resolution: Output resolution / aspect ratio label.
+        steps: Number of denoising steps (higher = more detail, slower).
+        guidance: Classifier-free guidance scale.
+        negative_prompt: What to avoid in the image (optional).
+        seed: Random seed for reproducibility.
+        randomize_seed: If true, pick a fresh random seed each run.
+
+    Returns:
+        The generated PNG image path and the seed that was used.
+    """
+    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))
+
+    result = pipe(
+        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",
+    )
+    img = result.video[0]
+    out_dir = Path(tempfile.mkdtemp(prefix="cosmos3_"))
+    p = out_dir / "image.png"
+    img.save(p)
+    return str(p), seed
+
+
+CSS = """
+.gradio-container { max-width: 1100px !important; margin: auto !important; }
+"""
+
+EXAMPLES = [
+    ["A photorealistic image of an autonomous delivery robot navigating a rainy city street at night, neon reflections on wet asphalt, cinematic lighting"],
+    ["A robotic arm on a factory assembly line precisely placing a component, industrial setting, sharp focus, high detail"],
+    ["A cozy reading nook with a robot sitting under a cherry blossom tree, holding an open book, soft afternoon light"],
+    ["An aerial view of a smart warehouse with automated guided vehicles moving between shelving, clean modern architecture"],
+]
+
+with gr.Blocks(theme=gr.themes.Soft(), css=CSS, title="Cosmos3-Super \u00b7 Text2Image") as demo:
+    gr.Markdown(
+        "# \U0001f30c Cosmos3-Super-Text2Image\n"
+        "Demo of [nvidia/Cosmos3-Super-Text2Image](https://huggingface.co/nvidia/Cosmos3-Super-Text2Image) "
+        "\u2014 a **64B**-parameter omnimodal world model for Physical AI \u2014 generating high-fidelity "
+        "images from text on a single Blackwell GPU via **NVFP4** weight-only quantization.\n\n"
+        "> NVIDIA officially tests this checkpoint only at BF16; NVFP4 is unofficial and may show "
+        "quality drift vs. the full-precision recipe."
+    )
+    with gr.Row():
+        prompt = gr.Textbox(
+            show_label=False,
+            placeholder="A photo of a robot reading a book under a cherry tree\u2026",
+            container=False,
+            scale=4,
+        )
+        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="1024\u00d71024 (1:1)",
+        )
+        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)
+        with gr.Row():
+            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+            seed = gr.Number(label="Seed", value=0, precision=0)
+
+    inputs = [prompt, resolution, steps, guidance, negative_prompt, seed, randomize_seed]
+    outputs = [out, seed]
+
+    gr.Examples(
+        examples=EXAMPLES,
+        inputs=[prompt],
+        outputs=outputs,
+        fn=generate,
+        cache_examples=True,
+        cache_mode="lazy",
+    )
+
+    run.click(generate, inputs, outputs)
+    prompt.submit(generate, inputs, outputs)
+
+demo.queue().launch(mcp_server=True)

requirements.txt

@@ -0,0 +1,10 @@
+diffusers @ git+https://github.com/huggingface/diffusers.git
+transformers
+accelerate
+torchvision
+torchaudio
+av
+imageio
+imageio-ffmpeg
+sentencepiece
+torchao