Add MiniMax-H3 prompt rewriter (Qwen3.6-27B + LoRA)

README.md

@@ -1,12 +1,59 @@
 ---
 title: MiniMax-H3 Prompt Rewriter
-emoji: ✏️
+emoji: 🎬
 colorFrom: red
 colorTo: pink
 sdk: gradio
 sdk_version: 6.26.0
 app_file: app.py
-short_description: Rewrite prompts into structured MiniMax-H3 video prompts
+short_description: Expand ideas into structured MiniMax-H3 A/V prompts
 python_version: "3.12"
-startup_duration_timeout: 30m
----
+startup_duration_timeout: 1h
+---
+
+# 🎬 MiniMax-H3 Prompt Rewriter
+
+Turn a one-line idea into the structured **text-to-audio-video (T2VA)** prompt that
+[MiniMax-H3](https://huggingface.co/MiniMaxAI/MiniMax-H3) expects.
+
+- **LoRA adapter** — [`lightx2v/MiniMax-H3-Prompt-Rewriter-LoRA`](https://huggingface.co/lightx2v/MiniMax-H3-Prompt-Rewriter-LoRA)
+- **Base model** — [`Qwen/Qwen3.6-27B`](https://huggingface.co/Qwen/Qwen3.6-27B) (bf16)
+- **Video generation** — [LightX2V](https://github.com/ModelTC/LightX2V) (separate step, not part of this Space)
+
+## What it does
+
+Given a short prompt plus a target duration and aspect ratio, the rewriter emits three
+fields, in this exact order:
+
+```text
+integrated_multimodal_description: [Shot 1] ...
+overall_soundscape: ...
+non_diegetic_music: ...
+```
+
+The rewrite expands shot structure, timing, composition, camera motion, physical action
+and continuity, and adds synchronized diegetic sound plus a non-diegetic score — while
+preserving the original intent. Feed the result to LightX2V with the same duration and
+aspect ratio to render the final MP4 (video + synchronized audio).
+
+## Faithfulness to the reference implementation
+
+This Space mirrors `infer.py` / `prompt_template.py` from the adapter repo: the same
+system prompt, the same `resolution / duration / original_prompt` user message, the
+Qwen3.6 chat template with thinking disabled, and the same decoding defaults
+(greedy, `repetition_penalty=1.05`; sampling uses `temperature=0.7`, `top_p=0.8`,
+`top_k=20`).
+
+The **Use the rewriter LoRA** checkbox under *Advanced settings* can be turned off to
+run the plain Qwen3.6-27B baseline — the `--base-only` comparison documented in the
+adapter's model card.
+
+## Notes
+
+- Runs on ZeroGPU with `size="xlarge"`: the 27B base model is ~56 GB in bf16, above the
+  48 GB half-card tier, so the full 96 GB card is required. No quantization is used, so
+  output matches the reference bf16 recipe.
+- Text-only. The adapter's current release does not consume images, video or audio
+  references (FL2VA / Ref2VA rewriting is on its roadmap).
+- Example prompts are the ones the authors showcase in the adapter's model card and
+  README.

app.py

@@ -1,501 +1,427 @@
-"""MiniMax-H3 Prompt Rewriter — rewrite raw prompts into structured MiniMax-H3
-audio-video generation prompts using a Qwen2.5-Omni-7B LoRA adapter.
-
-Supports T2AV, I2AV, L2AV, FL2AV, and Ref2AV task modes with optional
-reference images.  The output is a production-ready structured text prompt
-with the required sections (integrated_multimodal_description,
-overall_soundscape, non_diegetic_music for base tasks; six sections for
-Ref2AV).
-"""
+"""MiniMax-H3 T2VA Prompt Rewriter — Gradio Space.
 
-import os
-os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
-
-import spaces  # MUST come before torch / any CUDA-touching import
-import gc
-import math
-import re
-import tempfile
-from collections import Counter
-from contextlib import contextmanager
-from decimal import Decimal, ROUND_HALF_UP
-from pathlib import Path
-from typing import Any
-
-import torch
-import gradio as gr
-from peft import PeftModel
-from transformers import (
-    Qwen2_5OmniProcessor,
-    Qwen2_5OmniThinkerForConditionalGeneration,
-)
+Qwen3.6-27B + the `lightx2v/MiniMax-H3-Prompt-Rewriter-LoRA` PEFT adapter.
 
-from system_prompt import system_prompt_for_task
+Turns a short text prompt (plus a target duration and aspect ratio) into the
+structured, production-ready audio-video prompt that MiniMax-H3 expects:
 
-# ---------------------------------------------------------------------------
-# Constants
-# ---------------------------------------------------------------------------
+    integrated_multimodal_description: [Shot 1] ...
+    overall_soundscape: ...
+    non_diegetic_music: ...
 
-BASE_MODEL = "Qwen/Qwen2.5-Omni-7B"
-ADAPTER = "lightx2v/MiniMax-H3-Prompt-Rewriter-LoRA-Omni"
-RATIOS = ("adaptive", "21:9", "16:9", "4:3", "1:1", "3:4", "9:16")
-TASKS = ("T2AV", "I2AV", "L2AV", "FL2AV", "Ref2AV")
+Mirrors the reference implementation shipped in the adapter repo
+(`infer.py` + `prompt_template.py`): same system prompt, same user-message
+format, same chat template with thinking disabled, same decoding defaults.
+"""
 
-REFERENCE_LABEL_RE = re.compile(
-    r"<?\b(Picture|Video|Audio)\s+(\d+)\b>?", re.IGNORECASE
-)
-LABEL_PREFIX = {"image": "Picture", "video": "Video", "audio": "Audio"}
-
-EXPECTED_SECTIONS = {
-    "t2av": (
-        "integrated_multimodal_description:",
-        "overall_soundscape:",
-        "non_diegetic_music:",
-    ),
-    "i2av": (
-        "integrated_multimodal_description:",
-        "overall_soundscape:",
-        "non_diegetic_music:",
-    ),
-    "l2av": (
-        "integrated_multimodal_description:",
-        "overall_soundscape:",
-        "non_diegetic_music:",
-    ),
-    "fl2av": (
-        "integrated_multimodal_description:",
-        "overall_soundscape:",
-        "non_diegetic_music:",
-    ),
-    "ref2av": (
-        "subject_definitions:",
-        "summary:",
-        "retention_analysis:",
-        "detailed_description:",
-        "overall_soundscape:",
-        "non_diegetic_music:",
-    ),
-}
-
-IMAGE_MAX_PIXELS = 301056
-VIDEO_MAX_PIXELS = 100352
-VIDEO_FPS = 1.0
-MAX_NEW_TOKENS = 4096
-
-# ---------------------------------------------------------------------------
-# Helpers (ported from infer.py, adapted for Gradio)
-# ---------------------------------------------------------------------------
-
-
-def normalize_task(value: str) -> str:
-    aliases = {
-        "t2v": "t2av", "t2va": "t2av", "t2av": "t2av",
-        "i2v": "i2av", "i2va": "i2av", "i2av": "i2av",
-        "l2v": "l2av", "l2va": "l2av", "l2av": "l2av",
-        "fl2v": "fl2av", "fl2va": "fl2av", "fl2av": "fl2av",
-        "flf2v": "fl2av", "flf2va": "fl2av", "flf2av": "fl2av",
-        "ref2v": "ref2av", "ref2va": "ref2av", "ref2av": "ref2av",
-    }
-    normalized = value.strip().lower()
-    return aliases[normalized]
+import os
 
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
 
-def h3_effective_duration(requested_duration: int) -> tuple[int, float]:
-    """Map seconds to MiniMax-H3's legal ``17*n+5`` frame grid at 24 fps."""
-    frames = math.ceil((24 * requested_duration - 5) / 17) * 17 + 5
-    return frames, frames / 24.0
+import gc
+import json
+import time
+import contextlib
+import threading
 
+import spaces  # MUST come before torch / transformers / peft
 
-def format_duration(value: float) -> str:
-    return str(Decimal(str(value)).quantize(Decimal("0.01"), rounding=ROUND_HALF_UP))
+import torch
+import transformers
+import gradio as gr
+from huggingface_hub import hf_hub_download
+from safetensors.torch import load_file
+from peft import LoraConfig, PeftModel, set_peft_model_state_dict
+from transformers import AutoTokenizer, TextIteratorStreamer
+
+from prompt_template import SYSTEM_PROMPT, build_messages
+
+BASE_MODEL = "Qwen/Qwen3.6-27B"
+ADAPTER_REPO = "lightx2v/MiniMax-H3-Prompt-Rewriter-LoRA"
+
+# Aspect ratios / duration range accepted by the reference `infer.py`.
+RESOLUTIONS = ["16:9", "21:9", "4:3", "1:1", "3:4", "9:16"]
+MIN_DURATION, MAX_DURATION = 4, 15
+
+
+def _get_model_class():
+    """Same resolution order as the adapter repo's reference `infer.py`."""
+    for name in ("AutoModelForImageTextToText", "AutoModelForVision2Seq"):
+        model_class = getattr(transformers, name, None)
+        if model_class is not None:
+            return model_class
+    raise RuntimeError(
+        "A recent Transformers version with AutoModelForImageTextToText "
+        "support is required for Qwen3.6."
+    )
 
 
-def output_schema_ok(task: str, text: str) -> bool:
-    positions = [text.find(section) for section in EXPECTED_SECTIONS[task]]
-    return all(position >= 0 for position in positions) and positions == sorted(positions)
+print(f"[boot] transformers {transformers.__version__}, torch {torch.__version__}", flush=True)
 
+print(f"[boot] loading tokenizer from {BASE_MODEL} ...", flush=True)
+tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)
+if tokenizer.pad_token_id is None:
+    tokenizer.pad_token = tokenizer.eos_token
 
-def build_messages_from_inputs(
-    task: str,
-    prompt: str,
-    duration: int,
-    resolution: str,
-    image_paths: list[str | None],
-) -> list[dict[str, Any]]:
-    """Build the chat messages for the model, mirroring infer.py's build_messages."""
-
-    task_lower = normalize_task(task)
-    _, effective_duration = h3_effective_duration(duration)
-    formatted_duration = f"{format_duration(effective_duration)}s"
-
-    # Filter out None image paths
-    valid_images = [p for p in image_paths if p is not None and str(p).strip()]
-
-    # Build references list matching the task requirements
-    references = []
-    for idx, img_path in enumerate(valid_images, start=1):
-        references.append({
-            "order": idx,
-            "type": "image",
-            "label": f"<Picture {idx}>",
-            "path": str(img_path),
-        })
-
-    user_content: list[dict[str, Any]] = []
-    if task_lower == "ref2av":
-        user_content.append({"type": "text", "text": "Ordered MiniMax-H3 references:\n"})
-
-    for index, reference in enumerate(references, start=1):
-        label = reference["label"]
-        if task_lower == "i2av":
-            heading = f"{label} — exact first frame at 0.00 seconds:\n"
-        elif task_lower == "l2av":
-            heading = f"{label} — exact final frame at {formatted_duration}:\n"
-        elif task_lower == "fl2av" and index == 1:
-            heading = f"{label} — exact first frame at 0.00 seconds:\n"
-        elif task_lower == "fl2av":
-            heading = f"{label} — exact final frame at {formatted_duration}:\n"
-        else:
-            heading = f"{label}:\n"
-        user_content.append({"type": "text", "text": heading})
-        user_content.append(
-            {
-                "type": "image",
-                "image": reference["path"],
-                "max_pixels": IMAGE_MAX_PIXELS,
-            }
+# `<|im_end|>` and `<|endoftext|>` — both are terminators for Qwen3.6.
+EOS_IDS = sorted(
+    {
+        i
+        for i in (
+            tokenizer.eos_token_id,
+            tokenizer.convert_tokens_to_ids("<|im_end|>"),
+            tokenizer.convert_tokens_to_ids("<|endoftext|>"),
         )
-
-    user_content.append(
-        {
-            "type": "text",
-            "text": (
-                ("\n" if references else "")
-                + "Rewrite request:\n"
-                f"task: {task_lower.upper()}\n"
-                f"resolution: {resolution}\n"
-                f"effective_duration: {formatted_duration}\n"
-                f"raw_prompt: {prompt}"
-            ),
-        }
-    )
-    return [
-        {
-            "role": "system",
-            "content": [
-                {"type": "text", "text": system_prompt_for_task(task_lower)}
-            ],
-        },
-        {"role": "user", "content": user_content},
-    ]
-
-
-# ---------------------------------------------------------------------------
-# Model loading (module scope, eager .to("cuda"))
-# ---------------------------------------------------------------------------
-
-print("Loading Qwen2.5-Omni-7B processor...", flush=True)
-processor = Qwen2_5OmniProcessor.from_pretrained(
-    BASE_MODEL, trust_remote_code=False
+        if isinstance(i, int) and i >= 0
+    }
 )
-processor.image_processor.max_pixels = IMAGE_MAX_PIXELS
-processor.video_processor.max_pixels = VIDEO_MAX_PIXELS
-processor.tokenizer.padding_side = "right"
-if processor.tokenizer.pad_token_id is None:
-    processor.tokenizer.pad_token = processor.tokenizer.eos_token
-
-print("Loading Qwen2.5-Omni-7B thinker model...", flush=True)
-model = Qwen2_5OmniThinkerForConditionalGeneration.from_pretrained(
+print(f"[boot] eos ids: {EOS_IDS}", flush=True)
+
+print(f"[boot] loading base model {BASE_MODEL} (bf16, ~56 GB) ...", flush=True)
+_t0 = time.time()
+model = _get_model_class().from_pretrained(
     BASE_MODEL,
-    torch_dtype=torch.bfloat16,
-    attn_implementation="sdpa",
-    trust_remote_code=False,
+    dtype=torch.bfloat16,
     low_cpu_mem_usage=True,
+    trust_remote_code=True,
+    attn_implementation="sdpa",
 )
-
-print("Loading LoRA adapter...", flush=True)
-# Load adapter with torch_device="cpu" to avoid safetensors trying to use
-# CUDA in the main process (ZeroGPU has no GPU at module scope).
-model = PeftModel.from_pretrained(
-    model, ADAPTER, is_trainable=False, torch_device="cpu"
+print(f"[boot] base model loaded in {time.time() - _t0:.0f}s", flush=True)
+
+# --- LoRA ---------------------------------------------------------------
+# `PeftModel.from_pretrained` resolves its load device with `infer_device()`,
+# which returns "cuda" under the ZeroGPU torch hijack — at module scope there
+# is no real GPU, so it fails. Load the adapter state dict explicitly on CPU
+# and attach it by hand instead; the whole model is moved with a single
+# `.to("cuda")` at the end so ZeroGPU can pack it.
+print(f"[boot] loading LoRA adapter {ADAPTER_REPO} ...", flush=True)
+_t0 = time.time()
+with open(hf_hub_download(ADAPTER_REPO, "adapter_config.json")) as fh:
+    _adapter_cfg = json.load(fh)
+
+peft_config = LoraConfig(
+    r=_adapter_cfg["r"],
+    lora_alpha=_adapter_cfg["lora_alpha"],
+    lora_dropout=_adapter_cfg["lora_dropout"],
+    target_modules=_adapter_cfg["target_modules"],
+    bias=_adapter_cfg["bias"],
+    task_type=_adapter_cfg["task_type"],
+    inference_mode=True,
 )
-model = model.to("cuda")
-model.eval()
+model = PeftModel(model, peft_config)
 
-# Determine input device and context limit
-_input_embeddings = model.get_input_embeddings()
-if _input_embeddings is not None and _input_embeddings.weight.device.type != "meta":
-    input_device = _input_embeddings.weight.device
-else:
-    input_device = torch.device("cuda")
-
-_config = model.config
-_ctx_candidates = (
-    getattr(_config, "max_position_embeddings", None),
-    getattr(getattr(_config, "text_config", None), "max_position_embeddings", None),
-    getattr(getattr(_config, "thinker_config", None), "max_position_embeddings", None),
-    getattr(
-        getattr(getattr(_config, "thinker_config", None), "text_config", None),
-        "max_position_embeddings",
-        None,
-    ),
+_adapter_state = load_file(
+    hf_hub_download(ADAPTER_REPO, "adapter_model.safetensors"), device="cpu"
 )
-context_limit = 32768
-for v in _ctx_candidates:
-    if v is not None:
-        context_limit = int(v)
-        break
-
-print(f"Model loaded. Context limit: {context_limit}", flush=True)
-
-
-def _encode(messages: list[dict[str, Any]]) -> dict[str, Any]:
-    """Encode messages for the model, mirroring infer.py's _encode."""
-    encoded = super(Qwen2_5OmniProcessor, processor).apply_chat_template(
-        [messages],
-        tokenize=True,
-        add_generation_prompt=True,
-        return_dict=True,
-        return_tensors="pt",
-        text_kwargs={"padding": False},
-        images_kwargs={"max_pixels": IMAGE_MAX_PIXELS},
-        videos_kwargs={
-            "max_pixels": VIDEO_MAX_PIXELS,
-            "fps": VIDEO_FPS,
-            "use_audio_in_video": False,
-        },
-        video_fps=VIDEO_FPS,
-        load_audio_from_video=False,
-    )
-    encoded["use_audio_in_video"] = False
-    for key in ("pixel_values", "pixel_values_videos", "input_features"):
-        value = encoded.get(key)
-        if isinstance(value, torch.Tensor) and torch.is_floating_point(value):
-            encoded[key] = value.to(dtype=torch.bfloat16)
-    return dict(encoded)
+_load_result = set_peft_model_state_dict(model, _adapter_state)
+_unexpected = list(getattr(_load_result, "unexpected_keys", []) or [])
+_missing = [k for k in getattr(_load_result, "missing_keys", []) or [] if "lora_" in k]
+print(
+    f"[boot] adapter attached in {time.time() - _t0:.0f}s "
+    f"({len(_adapter_state)} tensors, missing_lora={len(_missing)}, unexpected={len(_unexpected)})",
+    flush=True,
+)
+if _unexpected:
+    print(f"[boot] WARNING unexpected adapter keys (first 5): {_unexpected[:5]}", flush=True)
+if _missing:
+    print(f"[boot] WARNING missing LoRA keys (first 5): {_missing[:5]}", flush=True)
 
+del _adapter_state
+gc.collect()
 
-# ---------------------------------------------------------------------------
-# Inference function
-# ---------------------------------------------------------------------------
+model.eval()
+model = model.to("cuda")
+print("[boot] model ready on cuda (ZeroGPU packed).", flush=True)
+
+
+def render_chat_prompt(prompt: str, resolution: str, duration: int) -> str:
+    """Render the training-time chat prompt with thinking mode disabled."""
+    messages = build_messages(prompt, resolution, duration)
+    base = dict(tokenize=False, add_generation_prompt=True)
+    try:
+        text = tokenizer.apply_chat_template(messages, enable_thinking=False, **base)
+    except TypeError:
+        text = tokenizer.apply_chat_template(
+            messages, chat_template_kwargs={"enable_thinking": False}, **base
+        )
+    # Belt-and-braces: if the template ignored `enable_thinking`, force the
+    # empty reasoning block ourselves so the model answers directly.
+    if text.endswith("<think>\n"):
+        text = text[: -len("<think>\n")] + "<think>\n\n</think>\n\n"
+    return text
 
 
-@spaces.GPU(duration=120)
+def _estimate_duration(
+    prompt="",
+    duration=10,
+    resolution="16:9",
+    use_lora=True,
+    max_new_tokens=1536,
+    greedy=True,
+    temperature=0.7,
+    top_p=0.8,
+    top_k=20,
+    repetition_penalty=1.05,
+    seed=42,
+    *args,
+    **kwargs,
+):
+    """Seconds of ZeroGPU time to reserve; scales with the token budget."""
+    try:
+        budget = int(max_new_tokens)
+    except (TypeError, ValueError):
+        budget = 1536
+    return int(min(300, 40 + budget / 8))
+
+
+@spaces.GPU(duration=_estimate_duration, size="xlarge")
 def rewrite_prompt(
-    task: str,
     prompt: str,
-    image1: str | None,
-    image2: str | None,
-    duration: int,
-    resolution: str,
-    seed: int,
-    progress: gr.Progress(track_tqdm=True),
-) -> str:
-    """Rewrite a raw prompt into a structured MiniMax-H3 video-generation prompt.
+    duration: int = 10,
+    resolution: str = "16:9",
+    use_lora: bool = True,
+    max_new_tokens: int = 1536,
+    greedy: bool = True,
+    temperature: float = 0.7,
+    top_p: float = 0.8,
+    top_k: int = 20,
+    repetition_penalty: float = 1.05,
+    seed: int = 42,
+):
+    """Rewrite a short prompt into a structured MiniMax-H3 audio-video prompt.
 
     Args:
-        task: Task mode — T2AV (text-only), I2AV (image as first frame),
-            L2AV (image as last frame), FL2AV (two images as first & last frames),
-            or Ref2AV (images as general references).
-        prompt: The raw user prompt describing the desired video.
-        image1: First reference image (required for I2AV, L2AV, FL2AV, Ref2AV).
-        image2: Second reference image (required for FL2AV, optional for Ref2AV).
-        duration: Target video duration in seconds (4–15).
-        resolution: Aspect ratio — adaptive, 21:9, 16:9, 4:3, 1:1, 3:4, or 9:16.
-        seed: RNG seed for reproducibility.
+        prompt: The short original text prompt to expand.
+        duration: Target clip length in seconds (4-15).
+        resolution: Target aspect ratio, one of 16:9, 21:9, 4:3, 1:1, 3:4, 9:16.
+        use_lora: Use the MiniMax-H3 rewriter LoRA; False runs the plain Qwen3.6-27B baseline.
+        max_new_tokens: Maximum number of tokens to generate.
+        greedy: Deterministic greedy decoding; False enables sampling.
+        temperature: Sampling temperature (ignored when greedy).
+        top_p: Nucleus sampling top-p (ignored when greedy).
+        top_k: Top-k sampling cutoff (ignored when greedy).
+        repetition_penalty: Repetition penalty applied during decoding.
+        seed: RNG seed used for sampling.
+
+    Returns:
+        The rewritten, structured MiniMax-H3 prompt and a short status line.
     """
-    task_lower = normalize_task(task)
-
-    # Gather image paths
-    image_paths = []
-    for img in (image1, image2):
-        if img is not None and str(img).strip():
-            image_paths.append(str(img))
-
-    # Validate inputs
-    if not prompt or not prompt.strip():
-        return "Error: Please provide a prompt."
-
-    task_img_counts = {"t2av": 0, "i2av": 1, "l2av": 1, "fl2av": 2}
-    if task_lower in task_img_counts:
-        expected = task_img_counts[task_lower]
-        if len(image_paths) != expected:
-            if expected == 0:
-                return f"Error: {task.upper()} does not use reference images."
-            return f"Error: {task.upper()} requires exactly {expected} image(s); got {len(image_paths)}."
-    elif task_lower == "ref2av":
-        if len(image_paths) < 1:
-            return "Error: Ref2AV requires at least one reference image."
-
-    # Build messages
-    messages = build_messages_from_inputs(
-        task_lower, prompt.strip(), duration, resolution, image_paths
-    )
-
-    # Encode
-    inputs = _encode(messages)
-    input_length = int(inputs["input_ids"].shape[1])
-    available = context_limit - input_length
-    if available <= 0:
-        return f"Error: encoded input ({input_length} tokens) exceeds model context ({context_limit})."
-
-    max_new = min(MAX_NEW_TOKENS, available)
-    inputs = {
-        key: value.to(input_device) if isinstance(value, torch.Tensor) else value
-        for key, value in inputs.items()
-    }
-
-    # Seed
-    torch.manual_seed(seed)
-    if torch.cuda.is_available():
-        torch.cuda.manual_seed_all(seed)
-
-    generation_kwargs: dict[str, Any] = {
-        "max_new_tokens": max_new,
-        "pad_token_id": processor.tokenizer.pad_token_id,
-        "eos_token_id": processor.tokenizer.eos_token_id,
-        "do_sample": False,
-    }
-
-    with torch.inference_mode():
-        output_ids = model.generate(**inputs, **generation_kwargs)
+    prompt = (prompt or "").strip()
+    if not prompt:
+        yield "", "⚠️ Enter a prompt first."
+        return
 
-    generated_ids = output_ids[0, input_length:]
-    text = processor.decode(
-        generated_ids,
-        skip_special_tokens=True,
-        clean_up_tokenization_spaces=False,
-    ).strip()
+    duration = int(max(MIN_DURATION, min(MAX_DURATION, int(duration))))
+    if resolution not in RESOLUTIONS:
+        resolution = "16:9"
+    max_new_tokens = int(max(128, min(4096, int(max_new_tokens))))
 
-    del inputs, output_ids, generated_ids
-    gc.collect()
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
+    transformers.set_seed(int(seed))
 
-    if not text:
-        return "Error: model generated empty text."
+    text = render_chat_prompt(prompt, resolution, duration)
+    inputs = tokenizer(text, return_tensors="pt", add_special_tokens=False)
+    inputs = {k: v.to("cuda") for k, v in inputs.items()}
+    n_prompt_tokens = int(inputs["input_ids"].shape[1])
 
-    # Append schema check note
-    schema_ok = output_schema_ok(task_lower, text)
-    if not schema_ok:
-        text += "\n\n⚠️ Note: Output may be missing some required sections."
-
-    return text
+    streamer = TextIteratorStreamer(
+        tokenizer, skip_prompt=True, skip_special_tokens=True, timeout=600
+    )
+    generation_kwargs = dict(
+        **inputs,
+        streamer=streamer,
+        max_new_tokens=max_new_tokens,
+        do_sample=not greedy,
+        repetition_penalty=float(repetition_penalty),
+        pad_token_id=tokenizer.pad_token_id,
+        eos_token_id=EOS_IDS,
+    )
+    if not greedy:
+        generation_kwargs.update(
+            temperature=float(temperature),
+            top_p=float(top_p),
+            top_k=int(top_k),
+        )
 
+    mode = "LoRA rewriter" if use_lora else "Qwen3.6-27B base (no LoRA)"
+    yield "", f"⏳ Generating with **{mode}** · {duration}s · {resolution} …"
+
+    error: list = []
+
+    def _run():
+        try:
+            with torch.inference_mode():
+                if use_lora:
+                    ctx = contextlib.nullcontext()
+                else:
+                    ctx = model.disable_adapter()
+                with ctx:
+                    model.generate(**generation_kwargs)
+        except Exception as exc:  # surfaced to the UI below
+            error.append(exc)
+            streamer.end()
+
+    started = time.perf_counter()
+    worker = threading.Thread(target=_run, daemon=True)
+    worker.start()
+
+    chunks = []
+    last_push = 0.0
+    for chunk in streamer:
+        chunks.append(chunk)
+        now = time.perf_counter()
+        if now - last_push > 0.25:
+            last_push = now
+            yield "".join(chunks).strip(), f"⏳ Generating with **{mode}** … {now - started:.0f}s"
+
+    worker.join()
+    if error:
+        yield "", f"❌ Generation failed: {error[0]}"
+        return
+
+    output = "".join(chunks).strip()
+    elapsed = time.perf_counter() - started
+    n_new = len(tokenizer(output, add_special_tokens=False)["input_ids"])
+    print(
+        f"[gen] mode={mode} prompt_tokens={n_prompt_tokens} new_tokens~{n_new} "
+        f"elapsed={elapsed:.1f}s ({n_new / max(elapsed, 1e-6):.1f} tok/s)",
+        flush=True,
+    )
+    yield output, (
+        f"✅ {mode} · {duration}s · {resolution} · ~{n_new} tokens in {elapsed:.1f}s"
+    )
 
-# ---------------------------------------------------------------------------
-# Gradio UI
-# ---------------------------------------------------------------------------
 
 CSS = """
 #col-container { max-width: 1100px; margin: 0 auto; }
 .dark .gradio-container { color: var(--body-text-color); }
 """
 
-with gr.Blocks() as demo:
-    gr.Markdown(
-        """
-        # MiniMax-H3 Prompt Rewriter
-        Rewrite raw prompts into structured **MiniMax-H3** audio-video generation prompts
-        using a Qwen2.5-Omni-7B LoRA adapter.
-
-        [Model card](https://huggingface.co/lightx2v/MiniMax-H3-Prompt-Rewriter-LoRA-Omni) ·
-        [Base model](https://huggingface.co/Qwen/Qwen2.5-Omni-7B)
-        """
-    )
+EXAMPLES = [
+    [
+        "Epic space-opera teaser: a captain watches the last fleet jump away, leaving her alone.",
+        10,
+        "16:9",
+    ],
+    ["A red fox walks through a snowy forest at dawn.", 15, "21:9"],
+    [
+        "ASMR close-up of a young woman's hands moving slowly in warm, soft light — "
+        "gentle tapping, sliding and fluttering against a blurred background.",
+        8,
+        "9:16",
+    ],
+    [
+        "A dynamic TV promo-style video of politicians at a neon-lit nighttime event, "
+        "shaking hands and exchanging smiles in a bustling modern room.",
+        12,
+        "16:9",
+    ],
+    [
+        "A man in casual attire reaches for a gaming joystick from a shelf in a well-lit "
+        "electronics store, shelves of consoles and accessories behind him.",
+        6,
+        "16:9",
+    ],
+]
+
+with gr.Blocks(theme=gr.themes.Citrus(), css=CSS, title="MiniMax-H3 Prompt Rewriter") as demo:
+    with gr.Column(elem_id="col-container"):
+        gr.Markdown(
+            "# 🎬 MiniMax-H3 Prompt Rewriter\n"
+            "Expand a one-line idea into the structured **text-to-audio-video** prompt "
+            "[MiniMax-H3](https://huggingface.co/MiniMaxAI/MiniMax-H3) expects — numbered shots, "
+            "camera motion, continuity, diegetic sound and score.\n\n"
+            "[`lightx2v/MiniMax-H3-Prompt-Rewriter-LoRA`](https://huggingface.co/lightx2v/MiniMax-H3-Prompt-Rewriter-LoRA) "
+            "(LoRA) on [`Qwen/Qwen3.6-27B`](https://huggingface.co/Qwen/Qwen3.6-27B) · "
+            "generate the actual video with [LightX2V](https://github.com/ModelTC/LightX2V)."
+        )
 
-    with gr.Row():
-        with gr.Column(scale=2):
-            task = gr.Dropdown(
-                choices=list(TASKS),
-                value="T2AV",
-                label="Task mode",
-                info="T2AV: text-only · I2AV: image→first frame · L2AV: image→last frame · FL2AV: two images→first+last · Ref2AV: images as references",
-            )
+        with gr.Row():
             prompt = gr.Textbox(
-                label="Raw prompt",
-                placeholder="Describe the video you want to generate…",
-                lines=5,
-            )
-            with gr.Row():
-                image1 = gr.Image(
-                    label="Reference image 1",
-                    type="filepath",
-                    visible=True,
-                )
-                image2 = gr.Image(
-                    label="Reference image 2",
-                    type="filepath",
-                    visible=True,
-                )
-            run_btn = gr.Button("Rewrite prompt", variant="primary", scale=1)
-
-        with gr.Column(scale=3):
-            output = gr.Textbox(
-                label="Structured MiniMax-H3 prompt",
-                lines=20,
+                label="Original prompt",
+                placeholder="A red fox walks through a snowy forest at dawn.",
+                lines=3,
+                scale=4,
             )
+            run_btn = gr.Button("Rewrite", variant="primary", scale=1)
 
-    with gr.Accordion("Advanced settings", open=False):
         with gr.Row():
             duration = gr.Slider(
-                minimum=4, maximum=15, value=5, step=1,
+                minimum=MIN_DURATION,
+                maximum=MAX_DURATION,
+                value=10,
+                step=1,
                 label="Duration (seconds)",
-                info="Target video duration (4–15 seconds)",
             )
             resolution = gr.Dropdown(
-                choices=list(RATIOS),
-                value="16:9",
-                label="Resolution / aspect ratio",
-            )
-            seed = gr.Number(
-                label="Seed", value=42, precision=0,
-                info="RNG seed for reproducibility",
+                choices=RESOLUTIONS, value="16:9", label="Aspect ratio"
             )
 
-    gr.Examples(
-        examples=[
-            # I2AV — one image as first frame
-            [
-                "I2AV",
-                "The video shows a young woman dressed in a gray hoodie as she stands in front of a bathroom mirror with a focus on grooming her hair. She has long, blonde hair cascading down her back. The bathroom is warmly lit, and a light is switched on above the mirror. The woman appears relaxed, and her hair is partially wet, suggesting she has recently taken a shower. The scene starts with her facing the mirror, her hair partially down and flowing. She lifts a section of her hair and secures it with a hair tie.",
-                "assets/examples/i2av/picture_1.jpg", None,
-                15, "adaptive", 42,
-            ],
-            # L2AV — one image as last frame
-            [
-                "L2AV",
-                "The video starts with a close-up shot that is severely out of focus, featuring primarily yellow, orange, and brown tones. As the video progresses, subtle shifts become apparent. By the final frames, details become even more discernible, with the presence of tiny blooming flowers now clearly visible. The focus improvement suggests a zooming mechanism or a subtle camera movement towards the subject matter, emphasizing the richness of the flora.",
-                "assets/examples/l2av/picture_1.jpg", None,
-                11, "adaptive", 42,
-            ],
-            # FL2AV — two images as first & last frames
-            [
-                "FL2AV",
-                "The video showcases a serene close-up of a green leaf, seemingly resting on foliage or vegetation, with water droplets scattered across its surface. The leaf appears lush and vibrant, displaying a deep green hue. The droplets vary in size and reflect light differently. The background is intentionally blurred, creating a depth-of-field effect. The scene does not exhibit any noticeable transitions; the frames are largely identical, suggesting either a very subtle movement within a single setting or the capturing of brief moments from a static position.",
-                "assets/examples/fl2av/picture_1.jpg", "assets/examples/fl2av/picture_2.jpg",
-                5, "adaptive", 42,
-            ],
-            # Ref2AV — two images as references
-            [
-                "Ref2AV",
-                "Create a 10-second video in 16:9 aspect ratio featuring a podcast recording session. Use <Picture 1> to define the host's appearance in an orange patterned shirt and the studio background with a green plant wall and blue accent lights. Use <Picture 2> to define the guest's appearance in a red shirt and their seating position opposite the host. The host speaks animatedly, gesturing with his hands over a zebra-print table, while the guest listens attentively. The lighting is warm and inviting. Audio should consist of clear dialogue between the two men, with minimal background noise.",
-                "assets/examples/ref2av/picture_1.jpg", "assets/examples/ref2av/picture_2.jpg",
-                10, "16:9", 42,
-            ],
-        ],
-        inputs=[task, prompt, image1, image2, duration, resolution, seed],
-        outputs=output,
-        fn=rewrite_prompt,
-        cache_examples=True,
-        cache_mode="lazy",
-    )
+        status = gr.Markdown("")
+        output = gr.Textbox(
+            label="Rewritten MiniMax-H3 prompt",
+            lines=22,
+            buttons=["copy"],
+        )
+
+        with gr.Accordion("Advanced settings", open=False):
+            use_lora = gr.Checkbox(
+                value=True,
+                label="Use the rewriter LoRA",
+                info="Uncheck to run the plain Qwen3.6-27B baseline for comparison.",
+            )
+            greedy = gr.Checkbox(
+                value=True, label="Greedy decoding", info="Deterministic; uncheck to sample."
+            )
+            max_new_tokens = gr.Slider(
+                minimum=256, maximum=4096, value=1536, step=128, label="Max new tokens"
+            )
+            temperature = gr.Slider(
+                minimum=0.0, maximum=2.0, value=0.7, step=0.05, label="Temperature (sampling)"
+            )
+            top_p = gr.Slider(minimum=0.05, maximum=1.0, value=0.8, step=0.05, label="Top-p")
+            top_k = gr.Slider(minimum=1, maximum=100, value=20, step=1, label="Top-k")
+            repetition_penalty = gr.Slider(
+                minimum=1.0, maximum=1.5, value=1.05, step=0.01, label="Repetition penalty"
+            )
+            seed = gr.Number(value=42, precision=0, label="Seed")
+
+        gr.Examples(
+            examples=EXAMPLES,
+            inputs=[prompt, duration, resolution],
+            outputs=[output, status],
+            fn=rewrite_prompt,
+            cache_examples=True,
+            cache_mode="lazy",
+        )
 
+        with gr.Accordion("System prompt used by the rewriter", open=False):
+            gr.Markdown(f"```text\n{SYSTEM_PROMPT}\n```")
+
+    all_inputs = [
+        prompt,
+        duration,
+        resolution,
+        use_lora,
+        max_new_tokens,
+        greedy,
+        temperature,
+        top_p,
+        top_k,
+        repetition_penalty,
+        seed,
+    ]
     run_btn.click(
         fn=rewrite_prompt,
-        inputs=[task, prompt, image1, image2, duration, resolution, seed],
-        outputs=output,
+        inputs=all_inputs,
+        outputs=[output, status],
         api_name="rewrite_prompt",
     )
+    prompt.submit(
+        fn=rewrite_prompt,
+        inputs=all_inputs,
+        outputs=[output, status],
+        api_name=False,
+    )
 
-demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)
+if __name__ == "__main__":
+    demo.launch(mcp_server=True)

prompt_template.py

@@ -0,0 +1,40 @@
+"""Prompt template shared by MiniMax-H3 T2VA prompt-rewriter inference."""
+
+from __future__ import annotations
+
+
+SYSTEM_PROMPT = """You are a professional prompt rewriter for joint audio-video generation.
+Rewrite the user's original prompt into one coherent, production-ready multimodal description for the requested output aspect ratio and duration.
+
+Return only these three fields, in this exact order:
+integrated_multimodal_description: ...
+overall_soundscape: ...
+non_diegetic_music: ...
+
+Requirements:
+- Expand the visual narrative into clearly numbered shots such as [Shot 1], [Shot 2], and include timestamps for cuts after the first shot when useful.
+- Make the number, timing, and pacing of shots appropriate for the requested duration.
+- Compose the scene for the requested aspect ratio.
+- Preserve the user's intent while adding concrete subjects, appearance, environment, lighting, composition, camera movement, physical motion, and temporal continuity.
+- Keep characters, objects, wardrobe, locations, and spatial relationships consistent across shots.
+- Describe synchronized diegetic audio in overall_soundscape and external score in non_diegetic_music.
+- Do not add explanations, Markdown fences, safety commentary, or fields other than the three requested fields."""
+
+
+def build_messages(prompt: str, resolution: str, duration: int) -> list[dict[str, str]]:
+    """Build the chat messages used during LoRA training and inference."""
+    prompt = prompt.strip()
+    if not prompt:
+        raise ValueError("prompt must not be empty")
+
+    return [
+        {"role": "system", "content": SYSTEM_PROMPT},
+        {
+            "role": "user",
+            "content": (
+                f"resolution: {resolution}\n"
+                f"duration: {duration}s\n"
+                f"original_prompt: {prompt}"
+            ),
+        },
+    ]

requirements.txt

@@ -1,10 +1,6 @@
-transformers>=5.0.0
-accelerate>=1.7.0
-peft>=0.15.2
-safetensors>=0.5.0
-qwen-omni-utils[decord]
-Pillow>=10.0.0
-librosa>=0.10.2
-soundfile>=0.12.1
-av>=12.0.0
-torchvision
+transformers>=5.16.0
+accelerate>=1.10
+peft>=0.20.0
+safetensors>=0.5
+torchvision
+pillow>=10