Blame for app.py

"""FrameForge demo — Gradio Space wrapping the official inference code.

Paper:  https://arxiv.org/abs/2610.36545
Code:   https://github.com/Lyne1/FrameForge (pinned below)
Weights: https://huggingface.co/lyneeeeeeeee/FrameForge (official LoRA)

Design notes:
- The official entry point (infer.py) rebuilds the whole pipeline per call.
  Here we build it ONCE at startup and reuse it per request, replicating the
  per-request tail of infer.py's generate_one_case() verbatim.
- CLI flags mirror scripts/infer.sh + configs/final.sh exactly.
"""

import os
import sys
import tarfile
import tempfile
import threading
import traceback
import urllib.request
from pathlib import Path

import spaces  # first: installs ZeroGPU CUDA emulation (effect-free off ZeroGPU)
import gradio as gr
import numpy as np
import torch
from huggingface_hub import hf_hub_download, snapshot_download

# --------------------------------------------------------------------------
# Paths and pinned versions
# --------------------------------------------------------------------------
FRAMEFORGE_COMMIT = "36d203cc49432e1cbe45ff0451b5e91157b550fb"  # Lyne1/FrameForge main
REPO_DIR = Path(os.environ.get("FF_REPO_DIR", "./frameforge_src"))
CKPT_DIR = Path(os.environ.get("FF_CKPT_DIR", "./ckpts"))
LORA_PATH = Path(os.environ.get("FF_LORA_PATH", "./checkpoints/pytorch_lora_weights.safetensors"))
HF_TOKEN = os.environ.get("HF_TOKEN")  # needed for gated FLUX.1-Redux-dev


def ensure_code() -> None:
    """Download the pinned FrameForge source tarball and put it on sys.path."""
    if not (REPO_DIR / "infer.py").exists():
        REPO_DIR.parent.mkdir(parents=True, exist_ok=True)
        url = f"https://codeload.github.com/Lyne1/FrameForge/tar.gz/{FRAMEFORGE_COMMIT}"
        print(f"Downloading FrameForge code from {url}")
        with urllib.request.urlopen(url) as resp:
            with tarfile.open(fileobj=resp, mode="r:gz") as tf:
                tf.extractall(REPO_DIR.parent)
        os.rename(REPO_DIR.parent / f"FrameForge-{FRAMEFORGE_COMMIT}", REPO_DIR)
    if str(REPO_DIR.resolve()) not in sys.path:
        sys.path.insert(0, str(REPO_DIR.resolve()))
    _patch_frameforge_source()


def _patch_frameforge_source() -> None:
    """Patch attention.py: flex_attention -> SDPA.

    flex_attention is torch.compile-wrapped; ZeroGPU does not support
    torch.compile (and torch 2.8.0 dynamo crashes on it). Their score_mod is
    exactly a boolean keep-mask (mask -> finfo.min), which is what SDPA's
    attn_mask does with fused kernels. Eager flex_attention would be a slow
    reference implementation by comparison. Fails loudly if the pinned
    source changes under us.
    """
    attn_file = REPO_DIR / "hyvideo/models/transformers/modules/attention.py"
    src = attn_file.read_text()
    old = "        hidden_states = flex_attention(query, key, value, score_mod=score_mod)"
    new = (
        "        if attn_mask is not None:\n"
        "            hidden_states = F.scaled_dot_product_attention(\n"
        "                query, key, value, attn_mask=attn_mask.unsqueeze(1))\n"
        "        else:\n"
        "            hidden_states = F.scaled_dot_product_attention(query, key, value)"
    )
    if old in src:
        attn_file.write_text(src.replace(old, new, 1))
        print("Patched attention.py: flex_attention -> SDPA (no torch.compile on ZeroGPU)")
    elif new.splitlines()[1].strip() in src:
        pass  # already patched
    else:
        raise RuntimeError(
            "FrameForge attention.py no longer contains the flex_attention anchor; "
            "update _patch_frameforge_source()."
        )


def ensure_checkpoints() -> None:
    """Download base model, encoders and the FrameForge LoRA.

    Layout matches the official README, minus byt5-small/Glyph-SDXL-v2:
    the final config (configs/final.sh) runs with disable_by5=True, so
    create_pipeline never loads them.
    """
    CKPT_DIR.mkdir(parents=True, exist_ok=True)
    if not (CKPT_DIR / "transformer/720p_t2v").exists():
        print("Downloading HunyuanVideo-1.5 (720p_t2v transformer + vae + scheduler)...")
        snapshot_download(
            "tencent/HunyuanVideo-1.5",
            local_dir=str(CKPT_DIR),
            allow_patterns=["transformer/720p_t2v/**", "vae/**", "scheduler/**", "config.json"],
            token=HF_TOKEN,
        )
    if not (CKPT_DIR / "text_encoder/llm").exists():
        print("Downloading Qwen2.5-VL-7B-Instruct text encoder...")
        snapshot_download("Qwen/Qwen2.5-VL-7B-Instruct",
                          local_dir=str(CKPT_DIR / "text_encoder/llm"), token=HF_TOKEN)
    if not (CKPT_DIR / "vision_encoder/siglip").exists():
        print("Downloading FLUX.1-Redux-dev vision encoder (GATED - needs HF_TOKEN with accepted license)...")
        snapshot_download("black-forest-labs/FLUX.1-Redux-dev",
                          local_dir=str(CKPT_DIR / "vision_encoder/siglip"), token=HF_TOKEN)
    if not LORA_PATH.exists():
        LORA_PATH.parent.mkdir(parents=True, exist_ok=True)
        print("Downloading FrameForge LoRA checkpoint...")
        hf_hub_download("lyneeeeeeeee/FrameForge", "pytorch_lora_weights.safetensors",
                        local_dir=str(LORA_PATH.parent), token=HF_TOKEN)


# --------------------------------------------------------------------------
# Pipeline setup (runs once)
# --------------------------------------------------------------------------
pipe = None
ff_args = None
ip = None
LOAD_ERROR = None
PIPE_LOCK = threading.Lock()


def _apply_zerogpu_shims() -> None:
    """Patches for code paths that require real CUDA at build time.

    ZeroGPU's emulation intercepts tensor .to("cuda") but NOT low-level calls
    like torch.cuda.set_device or get_device_properties. Both shims are
    behavior-preserving: world_size=1 makes process-group meshes degenerate
    (no collectives are ever issued), and get_gpu_memory only selects VAE
    tiling config (we report the 96 GB xlarge slice).
    """
    import hyvideo.commons as hv_commons
    import hyvideo.commons.parallel_states as hv_ps
    import hyvideo.pipelines.hunyuan_video_pipeline as hv_pipe

    class _TrivialMesh:
        """1x1 DeviceMesh stand-in; with world_size=1 it is never used."""

        def __init__(self, mesh_shape, mesh_dim_names):
            self.mesh_shape = tuple(mesh_shape)
            self.mesh_dim_names = tuple(mesh_dim_names or ())

        def __getitem__(self, _dim):
            return self

        def get_local_rank(self, *_a, **_k):
            return 0

        def get_group(self, *_a, **_k):
            return None

        def size(self, *_a, **_k):
            return 1

    _real_init_device_mesh = hv_ps.init_device_mesh

    def _init_device_mesh_safe(device_type, mesh_shape, mesh_dim_names=None, **kwargs):
        if all(int(d) == 1 for d in mesh_shape):
            return _TrivialMesh(mesh_shape, mesh_dim_names)
        return _real_init_device_mesh(
            device_type, mesh_shape, mesh_dim_names=mesh_dim_names, **kwargs
        )

    hv_ps.init_device_mesh = _init_device_mesh_safe

    def _fake_gpu_memory(*_a, **_k):
        return 96 * 1024**3  # ZeroGPU xlarge slice; selects the large-VAE-tile config

    hv_commons.get_gpu_memory = _fake_gpu_memory
    hv_pipe.get_gpu_memory = _fake_gpu_memory


def load_pipeline() -> None:
    global pipe, ff_args, ip
    import infer as ff_infer  # official entry point (repo root)
    import inference_utils as _ip
    ip = _ip
    _apply_zerogpu_shims()

    # Flag-for-flag parity with scripts/infer.sh + configs/final.sh.
    argv = [
        "--model_path", str(CKPT_DIR),
        "--resolution", "720p", "--reference_size", "512", "--video_length", "2",
        "--sr", "false", "--rewrite", "false", "--cfg_distilled", "false",
        "--enable_step_distill", "false", "--sparse_attn", "false",
        "--use_sageattn", "false", "--enable_cache", "false",
        "--offloading", os.environ.get("FF_OFFLOAD", "false"),
        "--group_offloading", "false", "--overlap_group_offloading", "false",
        "--save_pre_sr_video", "true",
        "--lora_rank", "32", "--lora_alpha", "32",
        "--target_name", "all-linear", "--modules_to_save", "gate",
        "--shared_seq", "true", "--num_experts", "3",
        "--use_reasoning_latent", "true", "--num_reason_token", "1",
        "--guidance_scale", "6.0",
        # placeholders, replaced per request:
        "--input_image", "placeholder.png", "--prompt", "placeholder",
        "--output_path", "placeholder.png",
        "--lora_path", str(LORA_PATH),
    ]
    ff_args = ff_infer.build_parser().parse_args(argv)

    # Same initialization order as infer.py main(), adapted for ZeroGPU:
    # - no torch.cuda.set_device: a ZeroGPU slice has one device, and the
    #   low-level call is not intercepted by the CUDA emulation at startup.
    # - pre-init a gloo (CPU) process group so their init_device_mesh("cuda")
    #   does not try to auto-init NCCL; with world_size=1 and sp=1 the meshes
    #   are 1x1, sp_group is None, and no collectives ever run.
    import torch.distributed as dist
    if not dist.is_initialized():
        dist.init_process_group(
            backend="gloo",
            init_method="tcp://127.0.0.1:29511",
            rank=0,
            world_size=1,
        )
    ip.initialize_parallel_state(sp=1)
    pipe, _ = ff_infer.build_pipeline(ff_args)
    print("FrameForge pipeline loaded.")


def _save_tmp_image(pil_image, suffix=".png") -> str:
    path = tempfile.NamedTemporaryFile(delete=False, suffix=suffix).name
    pil_image.convert("RGB").save(path)
    return path


def _generate_duration_seconds(*fn_args) -> int:
    """ZeroGPU dynamic duration; denoising steps is the last positional input."""
    steps = fn_args[-1] if fn_args else 50
    # headroom covers cold-worker weight streaming (~35 GB) on first request
    return int(90 + 2.5 * float(steps))


def _generate_impl(input_image, prompt, negative_prompt, guidance_scale, seed, steps):
    """Per-request logic replicating generate_one_case() from official infer.py."""
    if pipe is None:
        raise gr.Error(f"Pipeline failed to load at startup:\n{LOAD_ERROR}")
    if input_image is None:
        raise gr.Error("Upload an image first.")
    if not prompt or not prompt.strip():
        raise gr.Error("Prompt cannot be empty.")

    with PIPE_LOCK:
        ff_args.input_image = _save_tmp_image(input_image)
        ff_args.prompt = prompt
        ff_args.negative_prompt = negative_prompt or ""
        ff_args.guidance_scale = float(guidance_scale)
        ff_args.seed = int(seed)
        ff_args.num_inference_steps = int(steps)

        torch.manual_seed(ff_args.seed)
        torch.cuda.manual_seed_all(ff_args.seed)

        ref_latents = ip.ref_video_encoder(
            ff_args.input_image, pipe.vae, "cuda",
            ff_args.reference_size, ff_args.reference_size,
        )
        if ff_args.use_reasoning_latent:
            ref_latents = ip.append_reasoning_latents(ref_latents, ff_args.num_reason_token)

        with torch.no_grad():
            sample = pipe(
                enable_sr=ff_args.sr,
                prompt=ff_args.prompt,
                aspect_ratio=ff_args.aspect_ratio,
                num_inference_steps=ff_args.num_inference_steps,
                sr_num_inference_steps=None,
                video_length=ff_args.video_length,
                negative_prompt=ff_args.negative_prompt,
                seed=ff_args.seed,
                output_type="pt",
                prompt_rewrite=ff_args.rewrite,
                return_pre_sr_video=ff_args.save_pre_sr_video,
                generate_image=True,
                ref_latents=ref_latents,
                guidance_scale=ff_args.guidance_scale,
                # tiled VAE decode uses threaded CUDA work and can deadlock
                # in forked ZeroGPU workers; 512x512x1 frame fits VRAM easily
                enable_vae_tile_parallelism=False,
            ).videos[0]
        print("[frameforge] pipe() returned", flush=True)

        image = torch.from_numpy(np.array(sample.squeeze(1)))
        out_path = tempfile.NamedTemporaryFile(delete=False, suffix=".png").name
        ip.save_image(image, out_path)
        print(f"[frameforge] saved {out_path}", flush=True)
        return out_path


@spaces.GPU(size="xlarge", duration=_generate_duration_seconds)
def generate(input_image, prompt, negative_prompt, guidance_scale, seed, steps):
    return _generate_impl(input_image, prompt, negative_prompt, guidance_scale, seed, steps)


@spaces.GPU(size="xlarge", duration=_generate_duration_seconds)
def generate_canny(input_image, prompt, negative_prompt, auto_canny, guidance_scale, seed, steps):
    """Canny tab: optionally extract edges from a regular photo first."""
    if auto_canny and input_image is not None:
        import cv2
        from PIL import Image
        gray = cv2.cvtColor(np.array(input_image.convert("RGB")), cv2.COLOR_RGB2GRAY)
        edges = cv2.Canny(gray, 100, 200)
        input_image = Image.fromarray(edges).convert("RGB")
    return _generate_impl(input_image, prompt, negative_prompt, guidance_scale, seed, steps)


# --------------------------------------------------------------------------
# Startup: download code + checkpoints, then build the pipeline once.
# --------------------------------------------------------------------------
try:
    ensure_code()
    ensure_checkpoints()
    load_pipeline()
except Exception:
    LOAD_ERROR = traceback.format_exc()
    print(LOAD_ERROR)


# --------------------------------------------------------------------------
# Gradio UI
# --------------------------------------------------------------------------
DESCRIPTION = """
# 🖼️ FrameForge: Unified Image Generation from Video Model Priors

First public demo of **FrameForge** ([paper](https://arxiv.org/abs/2610.36545) ·
[code](https://github.com/Lyne1/FrameForge) ·
[weights](https://huggingface.co/lyneeeeeeeee/FrameForge)) — one model for
subject-driven generation, Canny-to-image, and instruction-guided editing,
built on HunyuanVideo 1.5. Outputs are 512×512; generation takes a couple of
minutes (50 denoising steps on a video-model backbone).
"""

with gr.Blocks(title="FrameForge") as demo:
    gr.Markdown(DESCRIPTION)
    if LOAD_ERROR is not None:
        gr.Markdown(f"⚠️ **Pipeline failed to load.** First run downloads ~40 GB of "
                    f"checkpoints — check the logs. Details:\n```\n{LOAD_ERROR[:3000]}\n```")

    with gr.Tab("✏️ Instruction-guided editing"):
        gr.Markdown("Upload a photo and describe the edit, e.g. *“Change the sky to sunset "
                    "while preserving the rest of the scene.”*")
        with gr.Row():
            with gr.Column():
                edit_in = gr.Image(label="Source image", type="pil")
                edit_prompt = gr.Textbox(label="Edit instruction", lines=2)
                edit_guidance = gr.Slider(label="Guidance scale", minimum=1.0, maximum=15.0,
                                          value=6.0, step=0.1)
                edit_neg = gr.Textbox(label="Negative prompt", value="", lines=1, visible=False)
                edit_seed = gr.Number(label="Seed", value=1, step=1)
                edit_steps = gr.Slider(label="Denoising steps", minimum=10, maximum=100, value=50, step=1)
                edit_btn = gr.Button("Edit image", variant="primary")
            edit_out = gr.Image(label="Result", type="filepath")
        edit_btn.click(generate, [edit_in, edit_prompt, edit_neg, edit_guidance, edit_seed, edit_steps], edit_out)

    with gr.Tab("🧍 Subject-driven generation"):
        gr.Markdown("Upload a subject photo and describe the new scene, e.g. *“the dog playing "
                    "on a beach.”*")
        with gr.Row():
            with gr.Column():
                subj_in = gr.Image(label="Subject image", type="pil")
                subj_prompt = gr.Textbox(label="Scene prompt", lines=2)
                subj_guidance = gr.Slider(label="Guidance scale", minimum=1.0, maximum=15.0,
                                          value=6.0, step=0.1)
                subj_neg = gr.Textbox(label="Negative prompt", value="", lines=1, visible=False)
                subj_seed = gr.Number(label="Seed", value=1, step=1)
                subj_steps = gr.Slider(label="Denoising steps", minimum=10, maximum=100, value=50, step=1)
                subj_btn = gr.Button("Generate", variant="primary")
            subj_out = gr.Image(label="Result", type="filepath")
        subj_btn.click(generate, [subj_in, subj_prompt, subj_neg, subj_guidance, subj_seed, subj_steps], subj_out)

    with gr.Tab("📐 Canny-to-image"):
        gr.Markdown("Upload a Canny edge map — or a regular photo with auto-extract enabled — "
                    "plus a prompt, e.g. *“A red car parked beside a lake at sunset.”*")
        with gr.Row():
            with gr.Column():
                canny_in = gr.Image(label="Edge map (or photo)", type="pil")
                canny_prompt = gr.Textbox(label="Prompt", lines=2)
                canny_auto = gr.Checkbox(label="Auto-extract Canny edges from my photo", value=False)
                canny_guidance = gr.Slider(label="Guidance scale", minimum=1.0, maximum=15.0,
                                           value=2.0, step=0.1)
                canny_neg = gr.Textbox(label="Negative prompt", value="", lines=1, visible=False)
                canny_seed = gr.Number(label="Seed", value=1, step=1)
                canny_steps = gr.Slider(label="Denoising steps", minimum=10, maximum=100, value=50, step=1)
                canny_btn = gr.Button("Generate from edges", variant="primary")
            canny_out = gr.Image(label="Result", type="filepath")
        canny_btn.click(generate_canny,
                        [canny_in, canny_prompt, canny_neg, canny_auto, canny_guidance, canny_seed, canny_steps],
                        canny_out)
        canny_example = REPO_DIR / "assets/canny.png"
        if canny_example.exists():
            gr.Examples(examples=[[str(canny_example),
                                   "A red car parked beside a lake at sunset",
                                   "", False, 2.0, 1, 50]],
                        inputs=[canny_in, canny_prompt, canny_neg, canny_auto,
                                canny_guidance, canny_seed, canny_steps],
                        outputs=canny_out, fn=generate_canny)

    gr.Markdown("""
---
*Unofficial demo built on the authors' released inference code and weights.
Model: [lyneeeeeeeee/FrameForge](https://huggingface.co/lyneeeeeeeee/FrameForge)
(LoRA over [HunyuanVideo 1.5](https://huggingface.co/tencent/HunyuanVideo-1.5), Tencent Hunyuan Community License).
If you use this in research, cite the [paper](https://arxiv.org/abs/2610.36545).*
""")

demo.queue(default_concurrency_limit=1).launch()