app.py

import os

# Allocator config for memory pressure (video DiTs hit transient allocation spikes)
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # MUST come before any torch / CUDA-touching import
import torch
import sys
import json
import tempfile
import time
from pathlib import Path
from huggingface_hub import hf_hub_download, snapshot_download

# ----------------------------------------------------------------------------
# Disable torch.compile at module scope — ZeroGPU doesn't support JIT compile
# The causal_model.py uses `torch.compile(flex_attention, ...)` at import time.
# We patch torch.compile to a no-op before importing the wan modules.
# ----------------------------------------------------------------------------
_original_torch_compile = torch.compile
torch.compile = lambda fn, *args, **kwargs: fn

# ----------------------------------------------------------------------------
# Download Wan2.1-T2V-1.3B base model and Echo-Infinity checkpoints
# ----------------------------------------------------------------------------
BASE_MODEL_DIR = Path("/home/user/wan_models/Wan2.1-T2V-1.3B")
BASE_MODEL_DIR.mkdir(parents=True, exist_ok=True)

print("[Echo-Infinity] Downloading Wan2.1-T2V-1.3B base model...")
snapshot_download(
    repo_id="Wan-AI/Wan2.1-T2V-1.3B",
    local_dir=str(BASE_MODEL_DIR),
    repo_type="model",
)
print("[Echo-Infinity] Wan2.1-T2V-1.3B downloaded.")

CKPT_DIR = Path("/home/user/checkpoints")
CKPT_DIR.mkdir(parents=True, exist_ok=True)

print("[Echo-Infinity] Downloading Echo-Infinity checkpoints...")
hf_hub_download(
    repo_id="Echo-Team/Echo-Infinity",
    filename="echo_infinity.pt",
    local_dir=str(CKPT_DIR),
    repo_type="model",
)
hf_hub_download(
    repo_id="Echo-Team/Echo-Infinity",
    filename="echo_infinity-long.pt",
    local_dir=str(CKPT_DIR),
    repo_type="model",
)
print("[Echo-Infinity] Checkpoints downloaded.")

# Make the project root importable
PROJECT_ROOT = Path("/home/user/app")
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

# Create symlink so wan_models/ resolves (wan_wrapper.py uses relative paths)
# wan_wrapper.py loads from wan_models/Wan2.1-T2V-1.3B/...
app_wan_models = PROJECT_ROOT / "wan_models"
if not app_wan_models.exists():
    try:
        app_wan_models.symlink_to(BASE_MODEL_DIR.parent, target_is_directory=True)
    except FileExistsError:
        pass

# ----------------------------------------------------------------------------
# Build the inference config (5s / 21 frames for ZeroGPU feasibility)
# ----------------------------------------------------------------------------
from omegaconf import OmegaConf

config = OmegaConf.create({
    "denoising_step_list": [1000, 750, 500, 250],
    "warp_denoising_step": True,
    "num_frame_per_block": 3,
    "model_name": "Wan2.1-T2V-1.3B",
    "model_kwargs": {
        "local_attn_size": 12,
        "timestep_shift": 5.0,
        "sink_size": 3,
    },
    "data_path": "inference/prompts/demo_5s.txt",
    "output_folder": "output/5s",
    "inference_iter": -1,
    "num_output_frames": 21,
    "use_ema": True,
    "seed": 0,
    "num_samples": 1,
    "save_with_index": True,
    "global_sink": True,
    "context_noise": 0,
    "generator_ckpt": str(CKPT_DIR / "echo_infinity.pt"),
    "memory_kwargs": {
        "enabled": True,
        "Q_frames": 3,
        "tokens_per_frame": 1560,
        "n_encoder_layers": 2,
        "hidden_dim": 1536,
        "num_heads": 12,
        "head_dim": 128,
        "initializer_range": 0.014,
        "rope": True,
        "qk_norm": True,
        "use_evicted_kv": False,
        "use_batch_update": False,
        "use_sink_anchor": False,
        "use_vib": False,
        "bptt_clips": 1,
        "gate_init_bias": 2.0,
        "encoder_lr_multiplier": 5.0,
        "normalize_memory_k": True,
    },
})

# ----------------------------------------------------------------------------
# Load the pipeline at module scope (ZeroGPU pattern)
# ----------------------------------------------------------------------------
from pipeline import CausalInferencePipeline
from utils.misc import set_seed

device = torch.device("cuda")
set_seed(config.seed)

print("[Echo-Infinity] Building CausalInferencePipeline...")
pipeline = CausalInferencePipeline(config, device=device)
pipeline.is_lora_enabled = False

# Load the Echo-Infinity checkpoint
print("[Echo-Infinity] Loading checkpoint...")
state_dict = torch.load(config.generator_ckpt, map_location="cpu", weights_only=False)
if "generator" in state_dict or "generator_ema" in state_dict:
    if config.use_ema and "generator_ema" in state_dict:
        raw_gen_state_dict = state_dict["generator_ema"]
        if "generator" in state_dict:
            enc_keys = {k: v for k, v in state_dict["generator"].items() if "query_memory_encoder" in k}
            if enc_keys:
                raw_gen_state_dict = dict(raw_gen_state_dict)
                raw_gen_state_dict.update(enc_keys)
    else:
        raw_gen_state_dict = state_dict.get("generator", state_dict.get("generator_ema"))
elif "model" in state_dict:
    raw_gen_state_dict = state_dict["model"]
else:
    raise ValueError(f"Generator state dict not found in {config.generator_ckpt}")

def _clean_key(name):
    return name.replace("_fsdp_wrapped_module.", "")

cleaned_state_dict = {_clean_key(k): v for k, v in raw_gen_state_dict.items()}
missing, unexpected = pipeline.generator.load_state_dict(cleaned_state_dict, strict=False)
if len(missing) > 0:
    print(f"[Warning] {len(missing)} parameters missing: {missing[:8]} ...")
if len(unexpected) > 0:
    print(f"[Warning] {len(unexpected)} unexpected parameters: {unexpected[:8]} ...")

pipeline = pipeline.to(dtype=torch.bfloat16)
pipeline.generator.to(device="cuda")
pipeline.vae.to(device="cuda")

# Restore torch.compile (in case anything needs it later)
torch.compile = _original_torch_compile

print("[Echo-Infinity] Pipeline ready.")

# ----------------------------------------------------------------------------
# Inference function
# ----------------------------------------------------------------------------
from einops import rearrange
import imageio
import numpy as np
import random
import gradio as gr


def _estimate_duration(prompt: str, num_frames: int = 21, *args, **kwargs):
    """Estimate GPU duration based on number of frames."""
    # 21 frames with 4 denoising steps takes ~60-90s on A10G
    base = 60
    return min(300, base + int(num_frames * 2))


@spaces.GPU(duration=_estimate_duration)
def generate(
    prompt: str,
    seed: int = 0,
    num_frames: int = 21,
    progress=gr.Progress(track_tqdm=True),
):
    """Generate a short video from a text prompt using Echo-Infinity.

    Echo-Infinity uses a learnable evolving memory mechanism on top of
    Wan2.1-T2V-1.3B to enable real-time, constant-cost generation of
    arbitrary-length videos. This demo generates short clips (capped at
    21 frames ≈ 1.3 seconds at 16fps) to fit within ZeroGPU limits.

    Args:
        prompt: Text description of the video to generate.
        seed: Random seed for reproducibility.
        num_frames: Number of frames to generate (must be divisible by 3).
    """
    if not prompt or not prompt.strip():
        return None, "Please enter a text prompt."

    seed = int(seed)
    num_frames = int(num_frames)
    if num_frames % 3 != 0:
        num_frames = 21  # default

    set_seed(seed)
    torch.set_grad_enabled(False)

    # Prepare noise
    sampled_noise = torch.randn(
        [1, num_frames, 16, 60, 104],
        device="cuda",
        dtype=torch.bfloat16,
    )

    try:
        video, latents = pipeline.inference(
            noise=sampled_noise,
            text_prompts=[prompt],
            return_latents=True,
            low_memory=False,
            profile=False,
        )
    except Exception as e:
        return None, f"Error during generation: {str(e)}"

    # Convert to video frames
    current_video = rearrange(video, "b t c h w -> b t h w c").cpu()
    video = 255.0 * current_video
    pipeline.vae.model.clear_cache()

    # Save to temp file using imageio (torchvision.io.write_video not available)
    frames = video[0].numpy().astype(np.uint8)  # [T, H, W, C]
    tmp = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
    tmp.close()
    writer = imageio.get_writer(tmp.name, fps=16, codec="libx264")
    for frame in frames:
        writer.append_data(frame)
    writer.close()

    return tmp.name, f"Generated {num_frames} frames from seed {seed}."


# ----------------------------------------------------------------------------
# Gradio UI
# ----------------------------------------------------------------------------
CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

EXAMPLE_PROMPTS = [
    "A playful golden retriever puppy running through a green meadow covered with wildflowers. The puppy has a joyful expression, with its tail wagging energetically. It is bounding through the grass, leaving a trail of small footprints behind. The sky is bright blue with fluffy white clouds, and the sun casts a warm glow over the scene.",
    "A person swimming in the ocean, with waves gently crashing around them. The swimmer is fully submerged, with only their head breaking the water's surface, taking deep breaths between strokes. The sun is setting, casting a warm golden glow over the water, and the sky is filled with hues of orange and pink.",
    "A black and white animation-style video of a panda drinking coffee in a cozy café in Paris. The panda is sitting at a small round table with a steaming cup of coffee in front of it. It holds the cup delicately with its paw, sipping slowly with a relaxed and content expression.",
]

with gr.Blocks() as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            "# Echo-Infinity: Real-Time Infinite Video Generation\n"
            "Learnable evolving memory for constant-cost video generation, "
            "built on Wan2.1-T2V-1.3B. Generate short video clips from text prompts.\n\n"
            "📖 [Paper](https://arxiv.org/abs/2606.04527) | "
            "🌐 [Project Page](https://echo-team-joy-future-academy-jd.github.io/Echo-Infinity/) | "
            "🤗 [Model](https://huggingface.co/Echo-Team/Echo-Infinity) | "
            "💻 [Code](https://github.com/Echo-Team-Joy-Future-Academy-JD/Echo-Infinity)"
        )

        with gr.Row():
            prompt = gr.Textbox(
                show_label=False,
                placeholder="Describe the video you want to generate...",
                container=False,
                scale=4,
            )
            run = gr.Button("Generate", variant="primary", scale=1)

        output_video = gr.Video(label="Generated Video")
        status = gr.Textbox(label="Status", interactive=False)

        with gr.Accordion("Advanced settings", open=False):
            seed = gr.Number(label="Seed", value=0, precision=0)
            randomize = gr.Checkbox(label="Randomize seed", value=True)
            num_frames = gr.Slider(
                label="Number of frames",
                minimum=3,
                maximum=33,
                step=3,
                value=21,
                info="Must be divisible by 3. Higher = longer video but more GPU time.",
            )

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

    def _run(prompt_val, seed_val, randomize_val, num_frames_val):
        if randomize_val:
            seed_val = random.randint(0, 2**31 - 1)
        video_path, status_msg = generate(prompt_val, seed_val, num_frames_val)
        return video_path, status_msg, gr.update(value=seed_val)

    run.click(
        fn=_run,
        inputs=[prompt, seed, randomize, num_frames],
        outputs=[output_video, status, seed],
    )

demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)