app.py

import os

# ZeroGPU: torch.compile / dynamo are unsupported — disable before torch import.
os.environ.setdefault("TORCH_COMPILE_DISABLE", "1")
os.environ.setdefault("TORCHDYNAMO_DISABLE", "1")

import random
import tempfile

import numpy as np
import spaces
import torch
import gradio as gr
from PIL import Image
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file

from diffusers import LTX2InContextPipeline
from diffusers.pipelines.ltx2.pipeline_ltx2_ic_lora import LTX2ReferenceCondition
from diffusers.pipelines.ltx2.utils import DISTILLED_SIGMA_VALUES
from diffusers.utils import encode_video

# --- Config -----------------------------------------------------------------
# FAST distilled variant of the reference-sheet IC-LoRA: 8-step schedule, CFG off.
# (The card's tuned recipe is non-distilled / 30 steps; this trades a little fidelity for speed.)
BASE_MODEL = "diffusers/LTX-2.3-Distilled-Diffusers"
LORA_REPO = "linoyts/LTX-2.3-loras"
LORA_FILE = "ltx-2.3-22b-ic-lora-ingredients-0.9"  # note: no .safetensors extension in the repo
LORA_SCALE = 1.4
FPS = 24
WIDTH, HEIGHT = 768, 448
NUM_FRAMES = 121
NUM_STEPS = len(DISTILLED_SIGMA_VALUES)  # 8
MAX_SEED = np.iinfo(np.int32).max
HF_TOKEN = os.environ.get("HF_TOKEN")

# --- Load pipeline once at module scope (ZeroGPU registers it) ---------------
pipe = LTX2InContextPipeline.from_pretrained(BASE_MODEL, torch_dtype=torch.bfloat16)
pipe.to("cuda")
pipe.vae.enable_tiling()

_lora_path = hf_hub_download(LORA_REPO, LORA_FILE, token=HF_TOKEN)
pipe.load_lora_weights(load_file(_lora_path), adapter_name="refsheet")
pipe.set_adapters("refsheet", LORA_SCALE)


# --- Helpers ----------------------------------------------------------------
def _build_prompt(sheet: str, action: str) -> str:
    return f"Reference sheet: {sheet.strip()}\n\nGenerated video: {action.strip()}"


def _duration(*args, **kwargs):
    return 200


# --- Inference --------------------------------------------------------------
@spaces.GPU(duration=_duration)
def generate(image, sheet, action, lora_scale, seed, randomize,
             progress=gr.Progress(track_tqdm=True)):
    if image is None:
        raise gr.Error("Please upload a reference sheet image.")
    if not sheet.strip():
        raise gr.Error("Describe the panels in the reference sheet (characters, props, location).")
    if not action.strip():
        raise gr.Error("Describe the action / shot you want generated.")

    if randomize:
        seed = random.randint(0, MAX_SEED)
    seed = int(seed)

    sheet_img = image.convert("RGB").resize((WIDTH, HEIGHT), Image.LANCZOS)
    ref = [sheet_img] * NUM_FRAMES

    pipe.set_adapters("refsheet", float(lora_scale))
    prompt = _build_prompt(sheet, action)
    ref_cond = LTX2ReferenceCondition(frames=ref, strength=1.0)

    video_out, _audio = pipe(
        prompt=prompt,
        negative_prompt="",
        reference_conditions=[ref_cond],
        reference_downscale_factor=1,
        width=WIDTH,
        height=HEIGHT,
        num_frames=NUM_FRAMES,
        frame_rate=FPS,
        num_inference_steps=NUM_STEPS,
        sigmas=DISTILLED_SIGMA_VALUES,
        guidance_scale=1.0,
        stg_scale=0.0,
        audio_guidance_scale=1.0,
        audio_stg_scale=0.0,
        generator=torch.Generator(device="cuda").manual_seed(seed),
        output_type="np",
        return_dict=False,
    )

    out_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
    encode_video(video_out[0], fps=FPS, output_path=out_path)
    return out_path, seed


# --- UI ---------------------------------------------------------------------
with gr.Blocks(title="LTX-2.3 Reference Sheet (Fast / Distilled)") as demo:
    gr.Markdown(
        "# ⚡ LTX-2.3 Reference-Sheet Control — Fast (Distilled)\n"
        "Same reference-sheet IC-LoRA, run on the **distilled** checkpoint with an 8-step schedule "
        "for fast generation. Supply a composite reference sheet (characters / props / location) and an "
        "action prompt. For maximum fidelity use the non-distilled demo (30 steps, guidance 4.0). "
        "IC-LoRA: [`linoyts/LTX-2.3-loras`](https://huggingface.co/linoyts/LTX-2.3-loras)."
    )
    with gr.Row():
        with gr.Column():
            image_in = gr.Image(type="pil", label="Reference sheet (one clean panel per element, black background)")
            sheet = gr.Textbox(
                label="Reference sheet description",
                placeholder="a young woman with red hair in a green jacket (face close-up + turnaround); a brass pocket watch; a cobblestone alley at night",
                lines=3,
            )
            action = gr.Textbox(
                label="Generated video (the action / shot)",
                placeholder="the woman walks down the alley and checks the pocket watch, slow dolly-in",
                lines=2,
            )
            with gr.Accordion("Settings", open=False):
                lora_scale = gr.Slider(0.8, 1.8, value=1.4, step=0.05, label="LoRA strength")
                randomize = gr.Checkbox(True, label="Randomize seed")
                seed = gr.Slider(0, MAX_SEED, value=42, step=1, label="Seed")
            run = gr.Button("Generate (fast)", variant="primary")
        with gr.Column():
            video_out = gr.Video(label="Generated video")
            used_seed = gr.Number(label="Seed used", interactive=False)

    run.click(
        generate,
        inputs=[image_in, sheet, action, lora_scale, seed, randomize],
        outputs=[video_out, used_seed],
    )

    gr.Examples(
        examples=[
            [
                "examples/sheet_camping.jpg",
                "a young child lying beside a golden retriever dog (left panel); a cluster of small camping tents in a green field (top right); misty green mountains over a calm lake (bottom right)",
                "the child and the golden retriever walk together across the grassy field toward the tents, gentle handheld camera, soft daylight",
                1.4, 42, False,
            ],
            [
                "examples/sheet_astronaut.jpg",
                "an astronaut in a white spacesuit holding a helmet (left panel); a vast misty mountain landscape over still water (right panel)",
                "the astronaut walks slowly across the misty shoreline looking around, slow dolly-in",
                1.4, 42, False,
            ],
            [
                "examples/sheet_woman_horse.jpg",
                "a smiling young woman with curly dark hair in a light top (left panel); a dappled grey horse standing on grass (top right); a green misty mountain meadow (bottom right)",
                "the woman walks up to the grey horse in the meadow and gently pets its neck, soft daylight",
                1.4, 42, False,
            ],
        ],
        inputs=[image_in, sheet, action, lora_scale, seed, randomize],
        outputs=[video_out, used_seed],
        fn=generate,
        cache_examples=True,
        cache_mode="lazy",
    )

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