app.py

"""Joy-LTX 2.5 - one prompt, one take, picture + sound (ZeroGPU demo).

JoyAI-Echo's performance on LTX-2.5: the Echo video attention / feed-forward delta on the official
LTX-2.5 dev transformer with the official distilled LoRA baked in at 0.5. This Space runs ComfyUI
headless with the ComfyUI-JoyLTX25 node pack and the comfy-native int8 checkpoint - the same graph as
the released JoyLTX25 canvases (the Multishot sampler with one shot: two passes, x2 latent upscale,
joint audio).

Everything GPU-side runs in a subprocess (ComfyUI's own venv) INSIDE the @spaces.GPU call, so the
Gradio process never touches CUDA. First call after a cold start builds the venv and pulls ~36 GB of
models; later calls only start ComfyUI (~40 s) + render.
"""
import os
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import json
import shutil
import subprocess
import sys
import time
import urllib.request
from pathlib import Path

import gradio as gr
import spaces
from huggingface_hub import hf_hub_download

ROOT = Path("/tmp/joyltx25")
COMFY = ROOT / "ComfyUI"
VENV = COMFY / ".venv"
VENV_PY = VENV / "bin" / "python"
PORT = 8199

DIT_REPO = "joeygambino/joyai-echo-ltx25-echoVid-comfy-native"
DIT_FILE = "LTX25dist-echoVid-070T30-v2-DiT-comfy-int8.safetensors"
DIT_100 = "LTX25dist-echoVid-100T50-v2-DiT-comfy-int8.safetensors"
LTX_REPO = "Lightricks/LTX-2.5"
VAE_V = "vae/ltx-2.5-video-vae-bf16.safetensors"
VAE_A = "vae/ltx-2.5-audio-vae-bf16.safetensors"
UPS = "latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors"
TENC = "text_encoders/gemma4-12b-with-proj-ltx-2.5-comfy-int8-convrot.safetensors"
PACK_GIT = "https://github.com/jlucasmcrell/ComfyUI-JoyLTX25.git"
COMFY_GIT = "https://github.com/comfyanonymous/ComfyUI.git"

TOKEN = os.environ.get("HF_TOKEN")
SETUP_OK, SETUP_ERR = True, ""

EXAMPLE = ("A woman in her thirties with dark hair tied back, in a rust-orange sweater, sits at a kitchen table "
           "late at night with a mug of coffee going cold, warm tungsten light, a dark window behind her. She looks "
           "at the camera and says, in a dry unhurried American voice, \"The whole street went dark for six minutes, "
           "and my lights got brighter.\" Her hands stay around the mug; a fridge hums; quiet room tone; no music. "
           "Static camera, one continuous take.")


def _run(cmd, cwd=None, check=True):
    print("[setup] $", cmd, flush=True)
    p = subprocess.run(cmd, shell=True, cwd=cwd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
    tail = p.stdout[-3000:] if p.stdout else ""
    if tail:
        print(tail, flush=True)
    if check and p.returncode != 0:
        raise RuntimeError(tail)
    return p.stdout


def _link(src, dst):
    dst = Path(dst); dst.parent.mkdir(parents=True, exist_ok=True)
    if dst.exists():
        return
    try:
        os.symlink(src, dst)
    except OSError:
        shutil.copy2(src, dst)


def setup():
    ROOT.mkdir(parents=True, exist_ok=True)
    if not (COMFY / "main.py").exists():
        _run(f"git clone --depth 1 {COMFY_GIT} ComfyUI", ROOT)
    pack = COMFY / "custom_nodes" / "ComfyUI-JoyLTX25"
    if not pack.exists():
        _run(f"git clone --depth 1 {PACK_GIT} ComfyUI-JoyLTX25", COMFY / "custom_nodes")
    if not VENV_PY.exists():
        print("[setup] building the ComfyUI venv (uv, torch cu128)...", flush=True)
        _run(f"{sys.executable} -m uv venv --python 3.12 {VENV}", COMFY)
        _run(f"{sys.executable} -m uv pip install --python {VENV_PY} --extra-index-url https://download.pytorch.org/whl/cu128 "
             f"torch torchvision torchaudio -r requirements.txt comfy-kitchen av", COMFY)
    models = COMFY / "models"
    dm = models / "diffusion_models"; dm.mkdir(parents=True, exist_ok=True)
    for f in (DIT_FILE, DIT_100):
        if not (dm / f).exists():
            print(f"[setup] downloading {f} ...", flush=True)
            _link(hf_hub_download(DIT_REPO, f, token=TOKEN), dm / f)
    for rel in (VAE_V, VAE_A, UPS, TENC):
        if not (models / rel).exists():
            print(f"[setup] downloading {rel} ...", flush=True)
            _link(hf_hub_download(LTX_REPO, rel, token=TOKEN), models / rel)
    print("[setup] ready", flush=True)


try:
    setup()
except Exception as e:  # the UI still loads and shows the error
    SETUP_OK, SETUP_ERR = False, str(e)[-1500:]
    print("[setup] FAILED:", SETUP_ERR, flush=True)


# --------------------------------------------------------------------------- graph
def graph(prompt, negative, width, height, frames, two_pass, seed, dose, video_cfg):
    dit = DIT_FILE if dose.startswith("070") else DIT_100
    return {
        "1": {"class_type": "UNETLoader", "inputs": {"unet_name": dit, "weight_dtype": "default"}},
        "2": {"class_type": "CLIPLoader", "inputs": {"clip_name": Path(TENC).name, "type": "ltxv", "device": "default"}},
        "3": {"class_type": "VAELoader", "inputs": {"vae_name": Path(VAE_V).name}},
        "4": {"class_type": "VAELoader", "inputs": {"vae_name": Path(VAE_A).name}},
        "5": {"class_type": "LatentUpscaleModelLoader", "inputs": {"model_name": Path(UPS).name}},
        "7": {"class_type": "JoyLTX_Multishot", "inputs": {
            "model": ["1", 0], "clip": ["2", 0], "video_vae": ["3", 0], "audio_vae": ["4", 0], "upscale_model": ["5", 0],
            "prompts": json.dumps({"prompts": [prompt]}), "negative": negative or "pc game, console game, video game, cartoon, childish, ugly",
            "width": int(width), "height": int(height), "frames_per_shot": int(frames), "shot_count": 1,
            "join": "fresh (independent shots)", "overlap": 3, "seed": int(seed), "seed_per_shot": True,
            "sampler_name": "euler_ancestral",
            "sigmas_pass1": "1.0, 0.99375, 0.9875, 0.98125, 0.975, 0.909375, 0.725, 0.421875, 0.0",
            "two_pass": bool(two_pass), "sigmas_pass2": "0.85, 0.7250, 0.4219, 0.0",
            "video_cfg": float(video_cfg), "audio_cfg": 1.0, "frame_rate": 24.0, "save_every_shot": False,
            "identity_ref": "off", "identity_strength": 0.0, "image_strength": 1.0, "keyframe_strength": 0.8, "ref_strength": 0.0}},
        "8": {"class_type": "CreateVideo", "inputs": {"images": ["7", 0], "audio": ["7", 1], "fps": 24.0}},
        "9": {"class_type": "SaveVideo", "inputs": {"video": ["8", 0], "filename_prefix": "video/joyltx25_space", "format": "auto", "codec": "auto"}},
    }


def _post(url, data=None):
    req = urllib.request.Request(url, data=json.dumps(data).encode() if data is not None else None,
                                 headers={"Content-Type": "application/json"}, method="POST" if data is not None else "GET")
    return json.loads(urllib.request.urlopen(req, timeout=60).read())


def _estimate(prompt, seconds=6, size="960 x 544 (16:9)", two_pass=True, *a, **k):
    frames = 1 + 8 * round((float(seconds) * 24) / 8)
    base = 150                         # ComfyUI start + model load
    per_frame = 1.6 if two_pass else 0.8
    return int(min(1500, base + frames * per_frame))


@spaces.GPU(duration=_estimate)
def generate(prompt, seconds=6, size="960 x 544 (16:9)", two_pass=True, dose="070T30 (default)", seed=553010,
             video_cfg=1.0, negative=""):
    if not SETUP_OK:
        raise gr.Error(f"Setup failed at startup: {SETUP_ERR}")
    if not str(prompt).strip():
        raise gr.Error("Write a prompt first.")
    w, h = [int(x) for x in size.split("(")[0].replace(" ", "").split("x")]
    frames = max(9, 1 + 8 * round((float(seconds) * 24) / 8))
    log = COMFY / "comfy_space.log"
    proc = subprocess.Popen([str(VENV_PY), "main.py", "--listen", "127.0.0.1", "--port", str(PORT),
                             "--disable-auto-launch", "--output-directory", str(COMFY / "output")],
                            cwd=COMFY, stdout=open(log, "w"), stderr=subprocess.STDOUT)
    try:
        t0 = time.time()
        while True:
            try:
                urllib.request.urlopen(f"http://127.0.0.1:{PORT}/system_stats", timeout=3); break
            except Exception:
                if proc.poll() is not None or time.time() - t0 > 240:
                    raise gr.Error("ComfyUI did not start:\n" + open(log, errors="ignore").read()[-2000:])
                time.sleep(2)
        g = graph(prompt, negative, w, h, frames, two_pass, seed, dose, video_cfg)
        r = _post(f"http://127.0.0.1:{PORT}/prompt", {"prompt": g, "client_id": "space"})
        if "prompt_id" not in r:
            raise gr.Error("ComfyUI rejected the graph:\n" + json.dumps(r)[:1500])
        pid = r["prompt_id"]
        while True:
            hist = _post(f"http://127.0.0.1:{PORT}/history/{pid}")
            if pid in hist:
                st = hist[pid]["status"]
                if st.get("status_str") == "error":
                    msgs = [m[1].get("exception_message", "") for m in st.get("messages", []) if m[0] == "execution_error"]
                    raise gr.Error("Render failed: " + (msgs[0][:1200] if msgs else open(log, errors="ignore").read()[-1500:]))
                outs = [o for v in hist[pid].get("outputs", {}).values() for o in v.get("images", []) + v.get("video", [])]
                if not outs:
                    raise gr.Error("No output file produced.")
                o = outs[0]
                path = COMFY / "output" / o.get("subfolder", "") / o["filename"]
                dst = Path("/tmp") / f"joyltx25_{int(time.time())}.mp4"
                shutil.copy2(path, dst)
                return str(dst), f"{w}x{h} · {frames} frames ({frames/24:.1f} s) · {'two-pass x2' if two_pass else 'single pass'} · {dose} · seed {seed} · {time.time()-t0:.0f} s incl. startup"
            if proc.poll() is not None:
                raise gr.Error("ComfyUI exited during the render:\n" + open(log, errors="ignore").read()[-2000:])
            time.sleep(3)
    finally:
        try:
            proc.terminate(); proc.wait(timeout=20)
        except Exception:
            proc.kill()


CSS = "#col-container {max-width: 1100px; margin: 0 auto;} video {max-height: 560px;}"

with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            "# 🎬 Joy-LTX 2.5 — one prompt, one take, picture + sound\n"
            "**JoyAI-Echo's performance on LTX-2.5's engine.** The Echo video attention / feed-forward delta on the official "
            "LTX-2.5 dev transformer with the official distilled LoRA baked in at 0.5 — few-step, joint audio + video, "
            "lip-sync and expression from Echo, any length on your own card. This demo renders one take with the same "
            "Multishot sampler the ComfyUI canvases use (two passes, x2 latent upscale).\n\n"
            "Models: [GGUF](https://huggingface.co/joeygambino/joyai-echo-ltx25-echoVid-gguf) (RTX 30/40) · "
            "[comfy-native](https://huggingface.co/joeygambino/joyai-echo-ltx25-echoVid-comfy-native) (RTX 50) · "
            "[DEV merges](https://huggingface.co/joeygambino/joyai-echo-ltx25-echoVid-dev) · "
            "Workflows: [ComfyUI-JoyLTX25](https://github.com/jlucasmcrell/ComfyUI-JoyLTX25). "
            "LTX-2.x Community License; generated content is machine-generated."
        )
        if not SETUP_OK:
            gr.Markdown(f"⚠️ **Startup setup failed:** `{SETUP_ERR}`")
        with gr.Row():
            with gr.Column(scale=3):
                prompt = gr.Textbox(label="Prompt (who, where, what they say in quotes, the sounds, the camera)", lines=7, value=EXAMPLE)
                with gr.Row():
                    seconds = gr.Slider(4, 12, value=6, step=1, label="seconds")
                    size = gr.Dropdown(["960 x 544 (16:9)", "544 x 960 (9:16)", "768 x 768 (1:1)"], value="960 x 544 (16:9)", label="render size (pass 1)")
                with gr.Row():
                    two_pass = gr.Checkbox(value=True, label="two-pass x2 (1920x1088 out)")
                    dose = gr.Dropdown(["070T30 (default)", "100T50 (livelier)"], value="070T30 (default)", label="dose")
                with gr.Row():
                    seed = gr.Number(value=553010, precision=0, label="seed")
                    video_cfg = gr.Slider(0.6, 1.2, value=1.0, step=0.05, label="video_cfg (1.0 = distilled; 0.7-0.85 = cooler grade, slower)")
                negative = gr.Textbox(label="negative (mostly inert at cfg 1.0)", value="")
                btn = gr.Button("Render", variant="primary")
            with gr.Column(scale=4):
                video = gr.Video(label="result (video + audio)", autoplay=True)
                info = gr.Markdown("")
        gr.Markdown("Tips: say the sounds you want (LTX invents drones for anything left unsaid), say the accent, keep the camera still or slow, "
                    "quote the line. Multishot (seamless joins, identity across cuts, refs by name) lives in the ComfyUI canvases.")
        btn.click(generate, [prompt, seconds, size, two_pass, dose, seed, video_cfg, negative], [video, info])

if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1).launch()