Upload folder using huggingface_hub

app.py

@@ -1,7 +1,7 @@
 """MiniMax-H3 video generation with synchronized audio — FL2VA (text / first-last-frame to video+audio).
 
-Based on the diffusers MiniMax-H3 pipeline. Uses Int8 weight-only quantization on both the
-33B transformer and the 32B Qwen3-VL text encoder to fit within ZeroGPU xlarge (96 GB VRAM).
+Uses the pruned NVFP4 transformer and local truncated NVFP4-AWQ Qwen3-VL conditioner for compact
+downloads (~42 GB total vs 119 GB BF16), fitting in ZeroGPU xlarge (96 GB VRAM).
 """
 
 from __future__ import annotations
@@ -66,17 +66,20 @@ def lower_duration_floor(seconds: float = MIN_UI_DURATION) -> None:
 
 
 PIPE = None
+COND_PIPE = None
+COND_ERROR: str | None = None
 LOAD_ERROR: str | None = None
 LOADED_IN: float | None = None
 
 
 def load_models() -> str | None:
-    """Load the full MiniMax-H3 FL2VA pipeline with Int8 quantization at startup.
+    """Load the compact generator and local truncated conditioner at startup.
 
-    Int8 weight-only quantization roughly halves VRAM for both the 33B transformer and
-    the 32B Qwen3-VL text encoder, bringing the total to ~63 GB — fits in xlarge (96 GB).
+    Uses MiniMaxH3GeneratorBlocks (only VAEs + schedulers + video_processor from MiniMaxAI/MiniMax-H3,
+    ~10 GB) plus the pruned NVFP4 transformer from lilcheaty/MiniMax-H3-NVFP4 (~16 GB) and the local
+    NVFP4-AWQ Qwen3-VL conditioner from Comfy-Org/MiniMax-H3 (~16 GB). Total download: ~42 GB.
     """
-    global PIPE, LOAD_ERROR, LOADED_IN
+    global PIPE, COND_PIPE, COND_ERROR, LOAD_ERROR, LOADED_IN
 
     if PIPE is not None or LOAD_ERROR is not None:
         return LOAD_ERROR
@@ -84,65 +87,49 @@ def load_models() -> str | None:
     started = time.time()
     try:
         import torch
-        from diffusers import ModularPipeline, MiniMaxH3Transformer3DModel, TorchAoConfig
-        from transformers import Qwen3VLForConditionalGeneration
-        from transformers import TorchAoConfig as TransformersTorchAoConfig
-        from torchao.quantization import Int8WeightOnlyConfig
+        from diffusers import ComponentsManager
 
-        lower_duration_floor()
+        from h3_split_blocks import MiniMaxH3GeneratorBlocks
 
-        print(f"[load] loading MiniMax-H3 from {MODEL_REPO} with Int8 quantization ...", flush=True)
-
-        pipe = ModularPipeline.from_pretrained(MODEL_REPO)
-
-        # Quantize the 33B transformer with Int8 weight-only quantization
-        pipe.update_components(
-            transformer=MiniMaxH3Transformer3DModel.from_pretrained(
-                MODEL_REPO,
-                subfolder="transformer",
-                dtype=torch.bfloat16,
-                quantization_config=TorchAoConfig(
-                    Int8WeightOnlyConfig(version=2),
-                    modules_to_not_convert=[
-                        "proj_in", "audio_proj_in", "context_embedder", "time_embedder", "time_proj",
-                        "token_refiner", "norm_out", "proj_out", "audio_proj_out",
-                    ],
-                ),
-            ),
-            # Quantize the 32B Qwen3-VL text encoder with Int8 weight-only quantization
-            text_encoder=Qwen3VLForConditionalGeneration.from_pretrained(
-                MODEL_REPO,
-                subfolder="text_encoder",
-                dtype=torch.bfloat16,
-                quantization_config=TransformersTorchAoConfig(
-                    Int8WeightOnlyConfig(version=2),
-                    modules_to_not_convert=[
-                        "model.visual",
-                        "model.language_model.embed_tokens",
-                        "model.language_model.norm",
-                        "lm_head",
-                    ],
-                ),
-            ),
+        lower_duration_floor()
+        manager = ComponentsManager()
+        blocks = MiniMaxH3GeneratorBlocks()
+        print(f"[gen] loading {[c.name for c in blocks.expected_components]} from {MODEL_REPO} ...", flush=True)
+        pipe = blocks.init_pipeline(MODEL_REPO, components_manager=manager, collection="h3")
+
+        # Load only VAEs, schedulers, and video_processor from the main repo (~10 GB)
+        pipe.load_components(
+            names=["vae", "audio_vae", "scheduler", "audio_scheduler", "video_processor"],
+            dtype=torch.bfloat16,
         )
-        pipe.load_components(dtype=torch.bfloat16)
 
-        # VAEs stay full precision — a bfloat16 audio VAE decodes the soundtrack ~20 dB too quiet.
-        # VAEs are small (~6 GB) so we can pack them at startup; the quantized transformer (~31 GB) and
-        # text encoder (~26 GB) are too large to pack alongside the BF16 download on disk, so they stay
-        # as CPU tensors and move to CUDA on the first GPU call.
-        pipe.vae.to("cuda")
-        pipe.audio_vae.to("cuda")
+        # Load the pruned NVFP4 transformer from the separate checkpoint repo
+        from h3_nvfp4 import load_transformer
+
+        pipe.update_components(transformer=load_transformer())
+        pipe.transformer.set_attention_backend("_native_cudnn")
 
-        # Use cuDNN fused attention (10-20% faster than SDPA, no extra deps)
+        # Load the local truncated NVFP4-AWQ Qwen3-VL conditioner
         try:
-            pipe.transformer.set_attention_backend("_native_cudnn")
-        except Exception:
-            pipe.transformer.set_attention_backend("sdpa")
+            from h3_local_conditioner import load_local_conditioner
+            from h3_split_blocks import MiniMaxH3ConditionerBlocks
+
+            print("[cond] loading the local truncated NVFP4-AWQ conditioner ...", flush=True)
+            text_encoder, tokenizer, processor = load_local_conditioner()
+            cond_pipe = MiniMaxH3ConditionerBlocks().init_pipeline(MODEL_REPO)
+            cond_pipe.update_components(
+                text_encoder=text_encoder,
+                tokenizer=tokenizer,
+                processor=processor,
+            )
+        except Exception as error:
+            traceback.print_exc()
+            COND_ERROR = f"{type(error).__name__}: {error}"
+            print(f"[cond] local load failed ({COND_ERROR})", flush=True)
 
-        PIPE = pipe
+        PIPE, COND_PIPE = pipe, cond_pipe
         LOADED_IN = time.time() - started
-        print(f"[load] ready in {LOADED_IN:.0f}s", flush=True)
+        print(f"[gen] ready in {LOADED_IN:.0f}s", flush=True)
     except Exception as error:
         traceback.print_exc()
         LOAD_ERROR = f"**Loading `{MODEL_REPO}` failed** after {time.time() - started:.0f}s: `{type(error).__name__}: {error}`"
@@ -153,14 +140,17 @@ def status() -> str:
     if LOAD_ERROR:
         return LOAD_ERROR
     if PIPE is None:
-        return f"Loading `{MODEL_REPO}` (~119 GB BF16, Int8-quantized at load). Watch the Space logs."
-    return f"Ready · Int8 quantized · loaded in {LOADED_IN:.0f}s"
+        return "Loading pruned NVFP4 transformer + local NVFP4 conditioner + full-precision VAEs (~42 GB). Watch the Space logs."
+    import h3_nvfp4
+    engine_status = h3_nvfp4.status()
+    cond_status = "local NVFP4-AWQ" if COND_PIPE is not None else f"unavailable ({COND_ERROR})"
+    return f"Ready · {engine_status} · VAEs full precision · loaded in {LOADED_IN:.0f}s · conditioner {cond_status}"
 
 
-# Duration estimation: linear in rows (matmuls) + quadratic (attention) + decode cost
-_DUR_B, _DUR_C = 1.5e-4, 5.0e-9
+# Duration estimation
+_DUR_B, _DUR_C = 1.1745e-4, 3.8396e-9
 _DECODE_BASE, _DECODE_PER_DEFAULT_CANVAS, _DEFAULT_CANVAS_PIXELS = 15, 15, 960 * 544 * 124
-_PAD = 15
+_PLACEMENT_ALLOWANCE, _PAD = 12, 10
 
 
 def get_duration(prompt, image, last_image, height, width, num_frames, steps, seed, *a, **k):
@@ -171,44 +161,54 @@ def get_duration(prompt, image, last_image, height, width, num_frames, steps, se
     rows = latent_frames * patches + (int(image is not None) + int(last_image is not None)) * patches
     denoise = steps * (_DUR_B * rows + _DUR_C * rows**2)
     decode = _DECODE_BASE + _DECODE_PER_DEFAULT_CANVAS * (height * width * num_frames) / _DEFAULT_CANVAS_PIXELS
-    # Extra allowance for quantized model being somewhat slower per step + cold-start weight transfer
-    quant_allowance = 30
-    return max(60, int(denoise + decode) + quant_allowance + _PAD)
+    local_conditioning = 20
+    return max(60, int(denoise + decode) + local_conditioning + _PLACEMENT_ALLOWANCE + _PAD)
 
 
 @spaces.GPU(duration=get_duration, size=GPU_SIZE)
-def generate(prompt, image, last_image, height, width, num_frames, steps, seed, progress=gr.Progress(track_tqdm=True)):
-    """Generate video with synchronized audio from text and optional keyframes.
-
-    Args:
-        prompt: Text description of the video to generate.
-        image: Optional first frame image (PIL Image).
-        last_image: Optional last frame image (PIL Image).
-        height: Output video height in pixels.
-        width: Output video width in pixels.
-        num_frames: Number of frames to generate (must be 17*n+5).
-        steps: Number of denoising steps.
-        seed: Random seed for reproducibility.
-    """
+def generate(prompt, image, last_image, height, width, num_frames, steps, seed,
+             acceleration="Balanced", progress=gr.Progress(track_tqdm=True)):
+    """Generate video with synchronized audio from text and optional keyframes."""
     import torch
-    from diffusers.utils import encode_video
 
-    # Move the quantized transformer and text encoder to CUDA on each cold worker.
-    # They were kept as CPU tensors at startup to avoid exceeding the 150 GB disk
-    # quota (BF16 download ~119 GB + packed copy would be too large).
-    PIPE.transformer.to("cuda")
-    PIPE.text_encoder.to("cuda")
+    # Move models to CUDA on each cold worker
+    if COND_PIPE is not None:
+        COND_PIPE.text_encoder.to("cuda")
+    PIPE.to("cuda")
 
-    state = PIPE(
+    # Local conditioning
+    condition_state = COND_PIPE(
         prompt=prompt,
         image=image,
         last_image=last_image,
         height=int(height),
         width=int(width),
-        num_frames=int(num_frames),
-        num_inference_steps=int(steps),
-        generator=torch.Generator("cpu").manual_seed(int(seed)),
     )
+    prompt_embeds = condition_state.get("prompt_embeds")
+    text_token_tags = condition_state.get("text_token_tags")
+
+    begin_request = getattr(PIPE.transformer, "begin_request", None)
+    end_request = getattr(PIPE.transformer, "end_request", None)
+    if begin_request is not None:
+        begin_request(int(steps), acceleration)
+    try:
+        with torch.inference_mode():
+            state = PIPE(
+                prompt_embeds=prompt_embeds.to("cuda", non_blocking=True),
+                text_token_tags=text_token_tags,
+                image=image,
+                last_image=last_image,
+                height=int(height),
+                width=int(width),
+                num_frames=int(num_frames),
+                num_inference_steps=int(steps),
+                generator=torch.Generator("cpu").manual_seed(int(seed)),
+            )
+    finally:
+        if end_request is not None:
+            end_request()
+
+    from diffusers.utils import encode_video
 
     videos = state.get("videos")
     audio = state.get("audio")
@@ -228,7 +228,8 @@ def generate(prompt, image, last_image, height, width, num_frames, steps, seed,
 
 
 def run_generate(prompt, image_path=None, last_image_path=None, canvas=DEFAULT_CANVAS,
-                 duration=5, steps=28, seed=42, progress=gr.Progress(track_tqdm=True)):
+                 duration=5, steps=28, seed=42, acceleration="Balanced",
+                 progress=gr.Progress(track_tqdm=True)):
     """Handle a generation request from the Gradio UI."""
     if LOAD_ERROR:
         raise gr.Error(LOAD_ERROR)
@@ -249,7 +250,7 @@ def run_generate(prompt, image_path=None, last_image_path=None, canvas=DEFAULT_C
     final_frame = keyframe(last_image_path) if last_image_path else None
 
     progress(0.05, desc=f"Denoising {steps} steps at {width}x{height}, {num_frames} frames ...")
-    path = generate(prompt, first_frame, final_frame, height, width, num_frames, steps, seed, progress)
+    path = generate(prompt, first_frame, final_frame, height, width, num_frames, steps, seed, acceleration, progress)
     return path
 
 
@@ -325,6 +326,12 @@ with gr.Blocks(title="MiniMax-H3") as demo:
             with gr.Accordion("Advanced options", open=False):
                 canvas = gr.Dropdown(label="Canvas", choices=list(CANVASES), value=DEFAULT_CANVAS)
                 duration = gr.Slider(label="Duration (s)", minimum=MIN_UI_DURATION, maximum=MAX_UI_DURATION, step=1, value=5)
+                acceleration = gr.Radio(
+                    label="Acceleration",
+                    choices=["Balanced", "Exact"],
+                    value="Balanced",
+                    info="Balanced uses adaptive step reuse; Exact evaluates every step.",
+                )
                 steps = gr.Slider(label="Steps", minimum=10, maximum=40, step=1, value=28)
                 seed = gr.Number(label="Seed", value=42, precision=0)
 
@@ -348,7 +355,7 @@ with gr.Blocks(title="MiniMax-H3") as demo:
 
     run.click(
         run_generate,
-        [prompt, image, last_image, canvas, duration, steps, seed],
+        [prompt, image, last_image, canvas, duration, steps, seed, acceleration],
         [video],
         api_name="generate",
     )

requirements.txt

@@ -7,8 +7,9 @@ torchvision==0.26.0
 # The Qwen3-VL processor decides the vision patch count, so a different minor changes the conditioning.
 transformers==5.8.0
 accelerate==1.14.0
-# Int8 weight-only quantization for the 33B transformer and 32B text encoder
-torchao==0.18.0
+# Blackwell-native NVFP4 GEMMs and the fused Q/K RMSNorm + split-half RoPE kernel used by h3_nvfp4.py.
+# CUDA 13 is mandatory: older builds emulate this path and are slower than BF16.
+comfy-kitchen==0.2.26
 # PyAV muxes the generated soundtrack onto the frames (encode_video)
 av
 pillow