app.py

"""
Alaya-EVOKE — text-to-video world model demo (t2v mode).

This Space drives the authors' own inference engine (scripts/inference/infer_single.py)
in-process, replicating the post-distillation t2v recipe exactly:
  * 3-step pyramid (STAGE2_STEPS = 1 1 1), no CFG (guidance_scale = 1.0)
  * 384x640 @ 24fps, no camera warp (t2v forbids warp — prompt only)

The heavy pipeline (14B DiT + UMT5 text encoder + Wan VAE) is loaded to GPU inside the
ZeroGPU worker and cached across requests by the engine's own _PIPE_CACHE.
"""
import os
import sys
import time
import glob
import shutil
import threading
import traceback
from pathlib import Path

import spaces
import gradio as gr
from huggingface_hub import snapshot_download

# ----------------------------------------------------------------------------------
# Paths / repo layout. The vendored engine resolves ckpt dirs relative to the repo root.
# ----------------------------------------------------------------------------------
REPO_ROOT = Path(__file__).resolve().parent
SCRIPTS_INFER = REPO_ROOT / "scripts" / "inference"
for _p in (str(REPO_ROOT), str(SCRIPTS_INFER)):
    if _p not in sys.path:
        sys.path.insert(0, _p)

MODEL_REPO = "SII-YuanyangYin/Evoke"
MODELS_DIR = REPO_ROOT / "models"
BASE_CKPT = MODELS_DIR / "evoke-base"
TRANSFORMER_PARENT = MODELS_DIR / "evoke" / "stage3_post_distillation"
OUT_ROOT = REPO_ROOT / "outputs"
OUT_ROOT.mkdir(parents=True, exist_ok=True)

HEIGHT, WIDTH, FPS = 384, 640, 24

# Authors' t2v negative prompt (scripts/inference/infer_batch.py NEG default).
NEG_PROMPT = (
    "oversaturated, garish colors, color shift, hue shift, color drift, inconsistent colors, "
    "color cast, color banding, flickering, jittery motion, abrupt transitions, sudden scene "
    "changes, temporal inconsistency, static, still picture, blurred details, subtitles, style, "
    "works, paintings, images, extra fingers, poorly drawn hands, poorly drawn faces, deformed, "
    "disfigured, misshapen limbs, fused fingers, three legs, many people in the background, "
    "walking backwards, messy background"
)


# ----------------------------------------------------------------------------------
# Weight download at module scope (no CUDA here). The Space downloads onto its own
# storage; only the two subtrees the t2v path needs are fetched.
# ----------------------------------------------------------------------------------
def _ensure_weights():
    MODELS_DIR.mkdir(parents=True, exist_ok=True)
    (MODELS_DIR / "evoke").mkdir(parents=True, exist_ok=True)
    need_base = not (BASE_CKPT / "model_index.json").is_file()
    need_tr = not (TRANSFORMER_PARENT / "transformer" / "config.json").is_file()
    if not (need_base or need_tr):
        print("[evoke] weights already present", flush=True)
        return
    print("[evoke] downloading weights (evoke-base + stage3_post_distillation) ...", flush=True)
    t0 = time.time()
    local = snapshot_download(
        MODEL_REPO,
        allow_patterns=[
            "evoke-base/**",
            "evoke/stage3_post_distillation/**",
        ],
    )
    local = Path(local)
    # Mirror the two subtrees into the repo-relative models/ layout the engine expects.
    for rel, dst in (
        ("evoke-base", BASE_CKPT),
        ("evoke/stage3_post_distillation", TRANSFORMER_PARENT),
    ):
        src = local / rel
        if src.is_dir() and not (dst / ("model_index.json" if rel == "evoke-base" else "transformer")).exists():
            dst.parent.mkdir(parents=True, exist_ok=True)
            if dst.exists() or dst.is_symlink():
                try:
                    dst.unlink()
                except IsADirectoryError:
                    shutil.rmtree(dst)
            os.symlink(src, dst)
    print(f"[evoke] weights ready in {time.time() - t0:.0f}s -> {MODELS_DIR}", flush=True)


_ensure_weights()

# Import the vendored engine only after sys.path is set. Module scope is torch/numpy only
# (no CUDA init, no depth-backend weight load), so this is safe outside @spaces.GPU.
import infer_single  # noqa: E402


def _build_argv(prompt: str, num_chunks: int, seed: int, out_dir: Path):
    """Reproduce the post_distill t2v recipe (scripts/inference/infer_post_distill.sh MODE=t2v).

    For t2v: GEO_ON=False -> no --use_geometric_state, no pose, no depth backend.
    IS_STAGE2=1, STAGE2_STEPS=1 1 1, GUIDANCE_SCALE=1.0, RESTRICT=0, NOISE_CENTER=0,
    VAE_DECODE_TYPE=persistent, WARP_SIGMA_MAX=0.135.
    """
    num_frames = 33 * int(num_chunks)  # NUM_CHUNKS -> NUM_FRAMES = 33 * chunks
    argv = [
        "--ckpt_path", str(BASE_CKPT),
        "--transformer_path", str(TRANSFORMER_PARENT),
        "--is_enable_stage2",
        "--stage2_num_stages", "3",
        "--stage2_steps", "1", "1", "1",
        "--stage2_stage_range", "0", "0.3333333333333333", "0.6666666666666666", "1",
        "--stage2_warp_compression_mode", "fixed_mem",
        "--height", str(HEIGHT), "--width", str(WIDTH),
        "--num_frames", str(num_frames), "--fps", str(FPS),
        "--num_inference_steps", "3",
        "--guidance_scale", "1.0",
        "--seed", str(int(seed)),
        "--vae_decode_type", "persistent",
        # Dump a self-contained mp4 for every chunk right after it decodes, so the UI can
        # stream the clip chunk-by-chunk (~1.5 s each) as it is produced instead of waiting
        # for the whole rollout. The engine still stitches the final geo_pred.mp4 at the end.
        "--save_chunk_segments",
        "--start_seconds", "0.0",
        "--ref_seconds", "0.0",
        "--no_raw_sink_frames",
        "--warp_noise_sigma_invisible", "1.0",
        "--warp_noise_sigma_min", "0.0",
        "--warp_noise_sigma_max", "0.135",
        "--joystick_hud", "auto",  # auto -> off for t2v (no pose)
        "--sample_type", "t2v",
        "--prompt", prompt,
        "--negative_prompt", NEG_PROMPT,
        "--image_noise_sigma_min", "0.0",
        "--image_noise_sigma_max", "0.0",
        "--output_folder", str(out_dir),
    ]
    return argv


def _estimate_duration(num_chunks: int) -> int:
    # Measured on xlarge ZeroGPU: 2 chunks incl. cold ~40 GB VRAM load ran in ~34 s.
    # Budget a cold-load base + per-chunk slope with a small safety margin so warm
    # calls stay lean without starving cold ones.
    return int(45 + 9 * int(num_chunks))


@spaces.GPU(duration=300, size="xlarge")
def _generate_gpu(prompt: str, num_chunks: int, seed: int):
    """Streaming generator: run the engine in a background thread and yield each per-chunk
    segment mp4 the moment it lands on disk, so the UI plays the video chunk-by-chunk as it
    is produced. The final full clip (geo_pred.mp4) is yielded once the rollout finishes.

    Yields (video_path, status_markdown) tuples.
    """
    out_dir = OUT_ROOT / f"t2v_{int(time.time() * 1000)}_{int(seed)}"
    if out_dir.exists():
        shutil.rmtree(out_dir, ignore_errors=True)
    out_dir.mkdir(parents=True, exist_ok=True)
    segments_dir = out_dir / "segments"

    argv = _build_argv(prompt, num_chunks, seed, out_dir)

    # Run the (blocking) engine in a daemon thread; ZeroGPU forks share threads. The thread
    # stashes any exception so the generator can re-raise it on the main (GPU) thread.
    _err = {}

    def _run():
        try:
            infer_single.main(argv)
        except BaseException as e:  # noqa: BLE001 - surface every engine failure to the UI
            _err["exc"] = e
            _err["tb"] = traceback.format_exc()

    t0 = time.time()
    worker = threading.Thread(target=_run, daemon=True)
    worker.start()

    def _segment(idx: int) -> Path:
        return segments_dir / f"segment_{idx:03d}_pred.mp4"

    yield None, "**Starting the world model…** (loading the 14B DiT — the first chunk takes longest)"

    # Stream segments as they complete. A segment file is only guaranteed fully flushed once
    # the *next* segment appears or the worker has exited, so we only yield segment N when
    # segment N+1 exists (or generation is done).
    n_expected = int(num_chunks)
    emitted = 0
    while True:
        worker_alive = worker.is_alive()
        # Emit every segment whose successor already exists (guaranteed fully written).
        while _segment(emitted + 1).is_file():
            seg = _segment(emitted)
            if seg.is_file():
                emitted += 1
                yield str(seg), (
                    f"**Streaming chunk {emitted} / {n_expected}** "
                    f"(~{emitted * 1.5:.1f} s generated so far)…"
                )
            else:
                emitted += 1
        if not worker_alive:
            break
        time.sleep(0.5)

    worker.join()
    if "exc" in _err:
        print(f"[evoke] engine failed:\n{_err.get('tb', '')}", flush=True)
        raise gr.Error("Generation failed (see Space logs for the engine trace).")

    elapsed = time.time() - t0
    print(f"[evoke] t2v done in {elapsed:.1f}s", flush=True)

    # Prefer the engine's final stitched clip; fall back to the last segment / any mp4.
    mp4 = out_dir / "geo_pred.mp4"
    if not mp4.is_file():
        cands = sorted(glob.glob(str(out_dir / "**" / "*.mp4"), recursive=True))
        if cands:
            mp4 = Path(cands[0])
        else:
            raise gr.Error("Generation produced no video (see Space logs for the engine trace).")
    yield str(mp4), f"**Done** — {n_expected} chunks (~{n_expected * 1.5:.1f} s) in {elapsed:.0f} s."


def generate(prompt, num_chunks=4, seed=44):
    if not prompt or not prompt.strip():
        raise gr.Error("Please enter a prompt.")
    num_chunks = int(num_chunks)
    # Re-point the ZeroGPU duration for this call based on the requested length.
    _generate_gpu.duration = _estimate_duration(num_chunks)
    # Forward the streamed (video, status) updates from the GPU worker to the UI.
    for video_path, status in _generate_gpu(prompt.strip(), num_chunks, int(seed)):
        yield video_path, status


# ----------------------------------------------------------------------------------
# UI
# ----------------------------------------------------------------------------------
DESCRIPTION = """
# Alaya-EVOKE · Text-to-Video World Model

Prompt-only **text-to-video** with the released **post-distillation** checkpoint of
[SII-YuanyangYin/Evoke](https://huggingface.co/SII-YuanyangYin/Evoke) — a 14B DiT world
model sampled in **3 steps, no CFG** at 384×640 · 24 fps.

This is the model's `t2v` mode: there is no camera control (the pipeline runs warp-free),
so the clip is generated from the prompt alone. One chunk ≈ 1.5 s of video, and the demo
**streams each chunk to the player as soon as it is decoded** — you watch the world roll out
chunk-by-chunk instead of waiting for the whole clip.
"""

EXAMPLE_PROMPTS = [
    ["A first-person walk down a narrow street at golden hour, low sun flaring between buildings, long shadows across the pavement."],
    ["A drone shot flying over a snowy mountain village at sunrise, smoke rising from chimneys."],
    ["Walking through a dense bamboo forest, dappled light shifting across the path, leaves swaying in the wind."],
    ["A first-person stroll along a tropical beach at sunset, gentle waves rolling onto the sand, palm trees on the right."],
]

CSS = """
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks(css=CSS, theme=gr.themes.Citrus(), title="Alaya-EVOKE Text-to-Video") as demo:
    gr.Markdown(DESCRIPTION)
    with gr.Row():
        with gr.Column(scale=3):
            prompt = gr.Textbox(
                label="Prompt",
                placeholder="Describe the scene and camera motion...",
                lines=3,
            )
            run = gr.Button("Generate video", variant="primary")
            with gr.Accordion("Advanced options", open=False):
                num_chunks = gr.Slider(
                    label="Length (chunks · 1 chunk ≈ 1.5 s)",
                    minimum=2, maximum=8, step=1, value=4,
                )
                seed = gr.Slider(
                    label="Seed", minimum=0, maximum=2**31 - 1, step=1, value=44,
                )
        with gr.Column(scale=4):
            out_video = gr.Video(label="Generated video (streams chunk-by-chunk)", autoplay=True)
            status = gr.Markdown("")

    gr.Examples(
        examples=EXAMPLE_PROMPTS,
        inputs=[prompt],
        outputs=[out_video, status],
        fn=generate,
        cache_examples=True,
        cache_mode="lazy",
    )

    run.click(fn=generate, inputs=[prompt, num_chunks, seed], outputs=[out_video, status])

if __name__ == "__main__":
    demo.queue(max_size=8).launch()