"""MiniMax H3 Character Swap LoRA — prompt + workflow demo.

This Space does not load the H3 base weights. It builds the Ref2VA prompt
and ComfyUI checklist from the published model card.
"""

from __future__ import annotations

import gradio as gr

LORA_REPO = "https://huggingface.co/akatz-ai/MiniMax-H3-Character-Swap-LoRA"
LORA_FILE = "h3_character_swap_pro4500_1000.safetensors"
DATASET = "https://huggingface.co/datasets/akatz-ai/H3-Character-Swap-v1"
BASE = "https://huggingface.co/Comfy-Org/MiniMax-H3"

DEFAULT_TARGET = "the man in the purple shirt"
DEFAULT_KEEP = (
    "Keep the replacement character's identity, outfit, and art style from <Picture 1>. "
    "Preserve the source video's camera, background, lighting, objects, and all other people. "
    "Match the target person's position, scale, pose, and movement. "
    "Do not show the reference sheet or its background."
)

SHORT_CAPTION = (
    "Swap {target} in <Video 1> with the character in <Picture 1>."
)

LONG_PROMPT = (
    "Replace only {target} in <Video 1> with the character in <Picture 1>. "
    "{keep}"
)


def build_pack(
    target: str,
    extra_keep: str,
    duration_s: float,
    fps: int,
    strength: float,
    use_long: bool,
) -> tuple[str, str, str]:
    """Build the Ref2VA prompt and a ComfyUI run sheet."""
    target = (target or DEFAULT_TARGET).strip()
    keep = (extra_keep or DEFAULT_KEEP).strip()
    prompt = LONG_PROMPT.format(target=target, keep=keep) if use_long else SHORT_CAPTION.format(target=target)

    sheet = f"""# H3 Character Swap — run sheet

## Weights
- LoRA: `{LORA_FILE}`
- Repo: {LORA_REPO}
- Strength: {strength:.2f}  (card default: 1.00)
- Trigger word: none
- Turbo / Spectrum / Sol: not required

## Base runtime (not bundled in this Space)
- Diffusion: `minimax_h3_ref2va_pruned_int8_convrot.safetensors`
- Source: {BASE}
- Video + audio VAEs from the same org pack
- Loader: model-only LoRA loader on the H3 transformer

## Inputs
- `<Video 1>` = source clip (prefer one continuous shot)
- `<Picture 1>` = replacement character or sheet (no sheet background in the output)
- Prompt below, verbatim

## Timing
- Duration request: {duration_s:.1f} s
- Frame rate: {fps} fps
- Card finding: 4–5 s continuous shots beat ~14 s tests
- Hard cuts often become zooms or slides — trim cuts first

## Limits (from the trainer)
- Background preservation is the advertised win vs base Ref2VA
- Expressions and cut timing are still unreliable
- Multi-character replacement was not supervised
- Audio from the model is not trustworthy; remux source audio after if needed
"""

    checklist = f"""1. Drop `{LORA_FILE}` into `ComfyUI/models/loras/`
2. Load H3 Ref2VA INT8 + VAEs (not this LoRA alone)
3. Apply LoRA at strength {strength:.2f} with a model-only loader
4. Wire source video → `<Video 1>`, character image → `<Picture 1>`
5. Paste the prompt
6. Generate a short 24 fps window, then continue if the join holds
7. If lips drift, copy original audio onto the new video in post
"""
    return prompt, sheet, checklist


def preview_note(video, image) -> str:
    v = "source video attached" if video else "no source video yet"
    i = "reference image attached" if image else "no reference image yet"
    return (
        f"Local preview only — this Space does **not** run MiniMax H3.\n"
        f"- {v}\n- {i}\n\n"
        f"Take both files into ComfyUI with the generated prompt."
    )


with gr.Blocks(title="MiniMax H3 Character Swap LoRA") as demo:
    gr.Markdown(
        f"""
# MiniMax H3 Character Swap LoRA
Experimental **one-character** replacement adapter by Akatz Labs
([model card]({LORA_REPO}), [dataset]({DATASET})).

Upload a clip and a character still, describe who to replace, and this demo
writes the Ref2VA prompt plus a ComfyUI checklist. **Live H3 inference is not
hosted here** — the adapter is 155 MB, the base Ref2VA stack is not, and no
Inference Provider serves it yet.
"""
    )

    with gr.Row():
        with gr.Column():
            video = gr.Video(label="Source clip → <Video 1>")
            image = gr.Image(type="filepath", label="Replacement character → <Picture 1>")
            target = gr.Textbox(
                label="Who to replace in the clip",
                value=DEFAULT_TARGET,
                placeholder="the woman in the red coat / the boy on the left / …",
            )
            extra_keep = gr.Textbox(
                label="Preservation instructions",
                value=DEFAULT_KEEP,
                lines=5,
            )
            with gr.Row():
                duration = gr.Slider(2, 14, value=5, step=0.5, label="Clip window (seconds)")
                fps = gr.Slider(8, 30, value=24, step=1, label="FPS")
                strength = gr.Slider(0.2, 1.2, value=1.0, step=0.05, label="LoRA strength")
            use_long = gr.Checkbox(
                value=True,
                label="Use the long preservation prompt (recommended in local evals)",
            )
            go = gr.Button("Build prompt + run sheet", variant="primary")
        with gr.Column():
            status = gr.Markdown("Attach a clip and a still, then build the pack.")
            prompt_out = gr.Textbox(label="Ref2VA prompt", lines=8)
            sheet_out = gr.Markdown()
            check_out = gr.Textbox(label="ComfyUI checklist", lines=10)

    go.click(
        preview_note,
        inputs=[video, image],
        outputs=status,
    ).then(
        build_pack,
        inputs=[target, extra_keep, duration, fps, strength, use_long],
        outputs=[prompt_out, sheet_out, check_out],
    )

    gr.Examples(
        examples=[
            [
                "the man in the purple shirt",
                DEFAULT_KEEP,
                5.0,
                24,
                1.0,
                True,
            ],
            [
                "the woman speaking on the left",
                DEFAULT_KEEP,
                4.0,
                24,
                1.0,
                True,
            ],
        ],
        inputs=[target, extra_keep, duration, fps, strength, use_long],
        label="Starter targets",
    )

    gr.Markdown(
        f"""
---
**Base model** is [Comfy-Org/MiniMax-H3]({BASE}), not merged into the LoRA.
License: MiniMax H3 Community License — read it before you ship anything.
"""
    )

if __name__ == "__main__":
    demo.launch(mcp_server=True)