import gradio as gr

MODEL_INFO = {
    "name": "MiniMax-H3-Fun-Controlnet-Union",
    "author": "alibaba-pai",
    "base_model": "MiniMaxAI/MiniMax-H3",
    "description": (
        "ControlNet-Union for MiniMax-H3 — a single checkpoint that conditions the "
        "MiniMax-H3 video generator on Canny, Depth, HED, MLSD, or Pose control videos, "
        "and also runs video inpainting. Trained with the VideoX-Fun pipeline."
    ),
}

RESULTS = [
    ("Canny", "A Tokyo street scene", "canny_tokyo_street"),
    ("Depth", "An astronaut floating in space", "depth_astronaut"),
    ("HED", "A T-Rex riding a BMX", "hed_trex_bmx"),
    ("MLSD", "A village with straight-line architecture", "mlsd_village"),
    ("Pose", "A dancer transformed into a flamenco dancer", "pose_dance"),
]

ASSET_URL = "https://huggingface.co/alibaba-pai/MiniMax-H3-Fun-Controlnet-Union/resolve/main"
RESULT_URL = "https://huggingface.co/alibaba-pai/MiniMax-H3-Fun-Controlnet-Union/resolve/main/results"


def result_section(control_type, caption, key):
    return gr.HTML(
        f"""
        <div style="display:flex; gap:12px; flex-wrap:wrap; margin:8px 0;">
          <div style="flex:1; min-width:240px;">
            <p style="font-weight:600; margin:0 0 4px;">Control — {control_type}</p>
            <video src="{ASSET_URL}/asset/{key}.mp4"
                   width="100%" controls muted loop
                   style="border-radius:8px; background:#111;"></video>
            <p style="font-size:0.8em; color:#888; margin:4px 0 0;">{caption}</p>
          </div>
          <div style="flex:1; min-width:240px;">
            <p style="font-weight:600; margin:0 0 4px;">Output — MiniMax-H3 + ControlNet-Union</p>
            <video src="{RESULT_URL}/{key}.mp4"
                   width="100%" controls muted loop
                   style="border-radius:8px; background:#111;"></video>
          </div>
        </div>
        """
    )


with gr.Blocks(
    theme=gr.themes.Soft(),
    title="MiniMax-H3-Fun-Controlnet-Union",
    css="""
        .model-header { text-align:center; padding:1.5rem 0 0.5rem; }
        .model-header h1 { font-size:1.8rem; font-weight:700; }
        .model-header .badge { display:inline-block; background:#6366f1; color:#fff;
            padding:0.2rem 0.7rem; border-radius:999px; font-size:0.75rem; margin-left:0.5rem; }
        video { box-shadow:0 4px 16px rgba(0,0,0,0.3); }
        .section-title { font-size:1.1rem; font-weight:600; margin-bottom:0.5rem;
            border-bottom:2px solid #e5e7eb; padding-bottom:0.3rem; }
        .info-table td { padding:0.4rem 0.8rem; vertical-align:top; }
        .info-table td:first-child { font-weight:600; white-space:nowrap; width:160px; }
    """
) as demo:
    gr.HTML(
        """
        <div class="model-header">
          <h1>MiniMax-H3-Fun-Controlnet-Union
            <span class="badge">alibaba-pai</span>
          </h1>
          <p style="color:#666; margin-top:0.3rem;">
            ControlNet-Union for <a href="https://huggingface.co/MiniMaxAI/MiniMax-H3" target="_blank">MiniMax-H3</a>
            &middot; <a href="https://github.com/aigc-apps/VideoX-Fun" target="_blank">VideoX-Fun</a> pipeline
          </p>
        </div>
        """
    )

    gr.Markdown("""
        **MiniMax-H3-Fun-Controlnet-Union** is a single ControlNet-Union checkpoint for the
        [MiniMax-H3](https://huggingface.co/MiniMaxAI/MiniMax-H3) video diffusion transformer.
        One model handles **Canny, Depth, HED, MLSD, and Pose** control conditions for
        video-to-video generation — no per-condition checkpoint switching — and also supports
        video inpainting.

        | File | Description |
        |------|-------------|
        | `MiniMax-H3-Fun-Controlnet-Union.safetensors` | Control branch weights (~6.8 GB): `control_proj_in` + 5 `control_blocks`. Loaded on top of the base MiniMax-H3 transformer. |
    """)

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### Model Features")
            gr.Markdown("""
                - **Union control** — one checkpoint for Canny, Depth, HED, MLSD, and Pose.
                - **5 control injection points** — layers 0, 10, 20, 30, 40 of the 50-block transformer.
                - **Guidance-distilled** — run with `guidance_scale = 1.0`; one forward pass per step.
                - **Inpainting** — control input widened to `control_in_dim = 49` (latent + masked latent + mask channels).
                - **`control_context_scale`** — scales every control skip before adding to the main branch:
                  `1.0` = strongest control, `0.0` = control branch off.
                - **Frame snap** — frame count snaps to the largest `17*n + 5` the video VAE can decode
                  (duration capped at 15 s), canvas keeps the control video's aspect ratio.
            """)
        with gr.Column(scale=1):
            gr.Markdown("### Inference Defaults")
            gr.HTML(
                """
                <table class="info-table">
                  <tr><td>num_inference_steps</td><td>40</td></tr>
                  <tr><td>guidance_scale</td><td>1.0 (guidance-distilled)</td></tr>
                  <tr><td>control_context_scale</td><td>1.00</td></tr>
                  <tr><td>seed</td><td>43</td></tr>
                  <tr><td>fps</td><td>24</td></tr>
                </table>
                """
            )

    gr.Markdown("---")
    gr.Markdown("### Results — All 5 Control Conditions")
    gr.Markdown(
        "All samples generated with `num_inference_steps=40`, `guidance_scale=1.0`, "
        "`control_context_scale=1.00`, seed 43."
    )

    for control_type, caption, key in RESULTS:
        result_section(control_type, caption, key)

    gr.Markdown("---")
    gr.Markdown("### How to Run Inference")
    gr.Markdown(
        """
        1. Clone the [VideoX-Fun](https://github.com/aigc-apps/VideoX-Fun) repository.
        2. Download the base **MiniMax-H3** model and this ControlNet-Union checkpoint.
        3. Place them under `models/Diffusion_Transformer/`:

        ```
        models/
        └── Diffusion_Transformer/
            ├── MiniMax-H3/                          # base transformer (~62 GB)
            └── MiniMax-H3-Fun-Controlnet-Union/
                └── MiniMax-H3-Fun-Controlnet-Union.safetensors   # control branch (~6.8 GB)
        ```

        4. Edit the variables at the top of `examples/minimax_h3_fun/predict_v2v_control.py`:

        ```python
        model_name          = "models/Diffusion_Transformer/MiniMax-H3"
        config_path         = "config/minimax_h3/minimax_h3_control.yaml"
        transformer_path    = "models/Diffusion_Transformer/MiniMax-H3-Fun-Controlnet-Union/MiniMax-H3-Fun-Controlnet-Union.safetensors"
        control_video       = "your_control_video.mp4"
        prompt              = "your prompt"
        ```

        5. Run: `python examples/minimax_h3_fun/predict_v2v_control.py`

        **Important notes:**
        - `config_path` must use the exact trained layout: `control_blocks_places: [0, 10, 20, 30, 40]`,
          `control_in_dim: 49`, `control_apply_audio: false`.
        - Keep `guidance_scale = 1.0` — higher values apply guidance twice and degrade output.
        - The control checkpoint carries **only the control branch**; the base MiniMax-H3 weights
          must be present at `model_name`.
        - **Memory:** transformer (~62 GB) + Qwen3-VL text encoder (~62 GB) do **not** fit a single
          80 GB GPU fully loaded. Use `model_group_offload` (fastest) or `model_cpu_offload_and_qfloat8`.
        """
    )

    gr.Markdown("---")
    gr.Markdown("### Links")
    gr.HTML(
        """
        <p>
          <a href="https://huggingface.co/alibaba-pai/MiniMax-H3-Fun-Controlnet-Union" target="_blank">
            Model page on Hugging Face
          </a>
          &middot;
          <a href="https://huggingface.co/MiniMaxAI/MiniMax-H3" target="_blank">
            MiniMax-H3 base model
          </a>
          &middot;
          <a href="https://github.com/aigc-apps/VideoX-Fun" target="_blank">
            VideoX-Fun repository
          </a>
        </p>
        """
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)