app.py

import os
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces
import sys
import tempfile
import math
import cv2
import numpy as np
import torch
import torchvision
import imageio
from PIL import Image
from omegaconf import OmegaConf
from einops import rearrange
from diffusers import FlowMatchEulerDiscreteScheduler

# Add current dir to path so videox_fun and vggt are importable
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

import gradio as gr

# ── Model paths ──────────────────────────────────────────────────────────
GEOWORLD_REPO = "peaes/GeoWorld"
WAN_FUN_REPO = "alibaba-pai/Wan2.1-Fun-1.3B-Control"
VGGT_REPO = "facebook/VGGT-1B"

WEIGHT_DTYPE = torch.bfloat16
CONFIG_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "wan_civitai.yaml")

# ── Download all weights via HF Hub at module scope ──────────────────────
from huggingface_hub import snapshot_download

print("Downloading Wan2.1-Fun-1.3B-Control base model...")
# Only download the files we need (skip optimizer states, etc.)
wan_fun_dir = snapshot_download(
    WAN_FUN_REPO, repo_type="model",
    allow_patterns=[
        "Wan2.1_VAE.pth",
        "config.json",
        "diffusion_pytorch_model.safetensors",
        "google/umt5-xxl/*",
        "models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth",
        "models_t5_umt5-xxl-enc-bf16.pth",
        "xlm-roberta-large/*",
    ],
)

print("Downloading GeoWorld stage1 & stage2 transformers...")
geoworld_dir = snapshot_download(
    GEOWORLD_REPO, repo_type="model",
    allow_patterns=[
        "geoworld_stage1/transformer/*",
        "geoworld_stage2/transformer/*",
    ],
)
geoworld_s1_dir = os.path.join(geoworld_dir, "geoworld_stage1", "transformer")
geoworld_s2_dir = os.path.join(geoworld_dir, "geoworld_stage2", "transformer")

print("Downloading VGGT-1B (safetensors only)...")
# Only download safetensors, not the .pt file (saves ~5 GB)
vggt_dir = snapshot_download(
    VGGT_REPO, repo_type="model",
    allow_patterns=["config.json", "model.safetensors"],
)

# ── Load models ──────────────────────────────────────────────────────────
from videox_fun.dist import set_multi_gpus_devices, shard_model
from videox_fun.models import (
    AutoencoderKLWan, AutoTokenizer, CLIPModel, WanT5EncoderModel,
)
from videox_fun.models import (
    WanTransformer3DModelGeoWorldstage1 as WanTransformer3DModelS1,
    WanTransformer3DModelGeoWorldstage2 as WanTransformer3DModelS2,
)
from videox_fun.models.cache_utils import get_teacache_coefficients
from videox_fun.pipeline import (
    WanFunControlPipelineGeoWorldstage1 as PipelineS1,
    WanFunControlPipelineGeoWorldstage2 as PipelineS2,
)
from videox_fun.utils.fp8_optimization import (
    convert_model_weight_to_float8, convert_weight_dtype_wrapper,
    replace_parameters_by_name,
)
from videox_fun.utils.utils import (
    filter_kwargs, get_image_latent, get_video_to_video_latent, save_videos_grid,
)
from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler

from vggt.models.vggt import VGGT


class VGGTWithDtype(VGGT):
    @property
    def dtype(self):
        return next(self.parameters()).dtype


device = "cuda"
config = OmegaConf.load(CONFIG_PATH)

# ── Stage 1: load transformer, vae, text encoder, tokenizer, clip, vggt ─
print("Loading Stage 1 transformer...")
transformer_s1 = WanTransformer3DModelS1.from_pretrained(
    geoworld_s1_dir,
    transformer_additional_kwargs=OmegaConf.to_container(
        config["transformer_additional_kwargs"]
    ),
    low_cpu_mem_usage=True,
    torch_dtype=WEIGHT_DTYPE,
)

print("Loading VAE...")
vae = AutoencoderKLWan.from_pretrained(
    os.path.join(wan_fun_dir, config["vae_kwargs"].get("vae_subpath", "vae")),
    additional_kwargs=OmegaConf.to_container(config["vae_kwargs"]),
).to(WEIGHT_DTYPE)

print("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(
    os.path.join(
        wan_fun_dir, config["text_encoder_kwargs"].get("tokenizer_subpath", "tokenizer")
    ),
)

print("Loading T5 text encoder...")
text_encoder = WanT5EncoderModel.from_pretrained(
    os.path.join(
        wan_fun_dir,
        config["text_encoder_kwargs"].get("text_encoder_subpath", "text_encoder"),
    ),
    additional_kwargs=OmegaConf.to_container(config["text_encoder_kwargs"]),
    low_cpu_mem_usage=True,
    torch_dtype=WEIGHT_DTYPE,
)
text_encoder = text_encoder.eval()

print("Loading CLIP image encoder...")
clip_image_encoder = CLIPModel.from_pretrained(
    os.path.join(
        wan_fun_dir,
        config["image_encoder_kwargs"].get("image_encoder_subpath", "image_encoder"),
    ),
).to(WEIGHT_DTYPE)
clip_image_encoder = clip_image_encoder.eval()

print("Loading VGGT...")
vggt = VGGTWithDtype.from_pretrained(vggt_dir).to(WEIGHT_DTYPE)
vggt = vggt.eval()

# ── Scheduler ───────────────────────────────────────────────────────────
sampler_name = "Flow"
shift = 3
Choosen_Scheduler = {
    "Flow": FlowMatchEulerDiscreteScheduler,
    "Flow_Unipc": FlowUniPCMultistepScheduler,
    "Flow_DPM++": FlowDPMSolverMultistepScheduler,
}[sampler_name]
scheduler = Choosen_Scheduler(
    **filter_kwargs(
        Choosen_Scheduler, OmegaConf.to_container(config["scheduler_kwargs"])
    )
)

# ── Build pipeline for Stage 1 ───────────────────────────────────────────
print("Building Stage 1 pipeline...")
pipeline_s1 = PipelineS1(
    transformer=transformer_s1,
    vae=vae,
    tokenizer=tokenizer,
    text_encoder=text_encoder,
    scheduler=scheduler,
    clip_image_encoder=clip_image_encoder,
    vggt=vggt,
).to(device)

# ── Load Stage 2 transformer ─────────────────────────────────────────────
print("Loading Stage 2 transformer...")
transformer_s2 = WanTransformer3DModelS2.from_pretrained(
    geoworld_s2_dir,
    transformer_additional_kwargs=OmegaConf.to_container(
        config["transformer_additional_kwargs"]
    ),
    low_cpu_mem_usage=True,
    torch_dtype=WEIGHT_DTYPE,
)

# ── Build pipeline for Stage 2 ───────────────────────────────────────────
scheduler2 = Choosen_Scheduler(
    **filter_kwargs(
        Choosen_Scheduler, OmegaConf.to_container(config["scheduler_kwargs"])
    )
)
print("Building Stage 2 pipeline...")
pipeline_s2 = PipelineS2(
    transformer=transformer_s2,
    vae=vae,
    tokenizer=tokenizer,
    text_encoder=text_encoder,
    scheduler=scheduler2,
    clip_image_encoder=clip_image_encoder,
    vggt=vggt,
).to(device)

# ── TeaCache setup ──────────────────────────────────────────────────────
enable_teacache = True
teacache_threshold = 0.10
num_skip_start_steps = 5
model_name_for_cache = "models/Diffusion_Transformer/Wan2.1-Fun-V1.1-1.3B-Control"
coefficients = get_teacache_coefficients(model_name_for_cache) if enable_teacache else None


def _enable_teacache(pipe, steps):
    if coefficients is not None:
        pipe.transformer.enable_teacache(
            coefficients, steps, teacache_threshold,
            num_skip_start_steps=num_skip_start_steps, offload=False,
        )


print("Models loaded successfully!")


@spaces.GPU(duration=300)
def generate_scene(
    input_image,
    partial_view_video,
    prompt: str = "",
    negative_prompt: str = "",
    seed: int = 43,
    guidance_scale: float = 6.0,
    num_inference_steps: int = 50,
    progress=gr.Progress(track_tqdm=True),
):
    """Generate a 3D scene video from an input image and a partial-view video.

    Runs the two-stage GeoWorld pipeline: Stage 1 uses the input image + partial-view
    video to produce an initial coarse 3D scene video. Stage 2 refines it using VGGT
    geometry features from the stage-1 output.

    Args:
        input_image: The reference/start image (PIL Image).
        partial_view_video: A partial-view video generated from the input image
            under a camera trajectory (video file path).
        prompt: Optional text prompt to guide generation.
        negative_prompt: Optional negative prompt.
        seed: Random seed for reproducibility.
        guidance_scale: CFG guidance scale.
        num_inference_steps: Number of diffusion steps.
    """
    if input_image is None:
        raise ValueError("Please provide an input image.")
    if partial_view_video is None:
        raise ValueError("Please provide a partial-view video.")

    sample_size = [480, 720]
    video_length = 49
    fps = 8

    # ── Stage 1 ──────────────────────────────────────────────────────────
    control_video_path = partial_view_video
    start_image_path = input_image  # could be a filepath or PIL image
    ref_image_path = input_image

    # Handle PIL Image input from Gradio
    if isinstance(start_image_path, Image.Image):
        tmp_img = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
        start_image_path.save(tmp_img.name)
        start_image_path = tmp_img.name
        ref_image_path = tmp_img.name

    clip_image = Image.open(ref_image_path).convert("RGB")
    ref_image_latent = get_image_latent(ref_image_path, sample_size=sample_size)
    start_image_latent = get_image_latent(start_image_path, sample_size=sample_size)

    # Read partial-view video for control
    input_video, input_video_mask, _, _ = get_video_to_video_latent(
        control_video_path, video_length=video_length, sample_size=sample_size,
        fps=fps, ref_image=None,
    )
    control_camera_video = None

    video_length_val = int(
        (video_length - 1) // vae.config.temporal_compression_ratio
        * vae.config.temporal_compression_ratio
    ) + 1

    generator = torch.Generator(device=device).manual_seed(seed)

    # Enable teacache for stage 1
    _enable_teacache(pipeline_s1, num_inference_steps)

    with torch.no_grad():
        sample_s1 = pipeline_s1(
            prompt,
            num_frames=video_length_val,
            negative_prompt=negative_prompt,
            height=sample_size[0],
            width=sample_size[1],
            generator=generator,
            guidance_scale=guidance_scale,
            num_inference_steps=num_inference_steps,
            control_video=input_video,
            control_camera_video=control_camera_video,
            ref_image=ref_image_latent,
            start_image=start_image_latent,
            clip_image=clip_image,
            shift=shift,
        ).videos

    # Save stage 1 output
    tmp_dir = tempfile.mkdtemp()
    stage1_video_path = os.path.join(tmp_dir, "stage1_output.mp4")
    save_videos_grid(sample_s1.cpu(), stage1_video_path, fps=fps)

    # ── Stage 2 ──────────────────────────────────────────────────────────
    # Read stage1 output video frames for VGGT
    vggt_video = []
    cap = cv2.VideoCapture(stage1_video_path)
    if not cap.isOpened():
        raise RuntimeError(f"Could not open stage1 output video: {stage1_video_path}")
    while True:
        ret, frame = cap.read()
        if frame is None:
            break
        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        img = Image.fromarray(frame_rgb).convert("RGB")
        vggt_video.append(img)
        if not ret:
            break
    cap.release()

    # Stage 2 control video = stage1 output
    input_video_s2, input_video_mask_s2, _, _ = get_video_to_video_latent(
        stage1_video_path, video_length=video_length_val, sample_size=sample_size,
        fps=fps, ref_image=None,
    )

    generator2 = torch.Generator(device=device).manual_seed(seed)

    # Enable teacache for stage 2
    _enable_teacache(pipeline_s2, num_inference_steps)

    clip_image_s2 = Image.open(ref_image_path).convert("RGB")
    ref_image_latent_s2 = get_image_latent(ref_image_path, sample_size=sample_size)
    start_image_latent_s2 = get_image_latent(start_image_path, sample_size=sample_size)

    with torch.no_grad():
        sample_s2 = pipeline_s2(
            prompt,
            num_frames=video_length_val,
            negative_prompt=negative_prompt,
            height=sample_size[0],
            width=sample_size[1],
            generator=generator2,
            guidance_scale=guidance_scale,
            num_inference_steps=num_inference_steps,
            control_video=input_video_s2,
            control_camera_video=None,
            ref_image=ref_image_latent_s2,
            start_image=start_image_latent_s2,
            clip_image=clip_image_s2,
            shift=shift,
            vggt_video=vggt_video,
        ).videos

    # Save stage 2 output
    stage2_video_path = os.path.join(tmp_dir, "stage2_output.mp4")
    save_videos_grid(sample_s2.cpu(), stage2_video_path, fps=fps)

    # Clean up temp image
    if isinstance(input_image, str) and os.path.exists(input_image):
        pass  # it was a user-provided file
    elif "tmp_img" in dir():
        os.unlink(tmp_img.name)

    return stage2_video_path


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

with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            "# GeoWorld: 3D Scene Generation from a Single Image\n"
            "Upload an image and a partial-view video to generate a high-fidelity "
            "3D scene video using the two-stage GeoWorld pipeline.\n\n"
            "[Paper](https://arxiv.org/abs/2511.23191) | "
            "[GitHub](https://github.com/peaes/GeoWorld) | "
            "[Model](https://huggingface.co/peaes/GeoWorld)"
        )

        with gr.Row():
            with gr.Column(scale=1):
                input_image = gr.Image(
                    label="Input Image",
                    type="filepath",
                    height=300,
                )
                partial_view_video = gr.Video(
                    label="Partial-View Video (camera trajectory)",
                    height=300,
                )
                run_btn = gr.Button("Generate 3D Scene", variant="primary")

            with gr.Column(scale=1):
                output_video = gr.Video(label="Generated 3D Scene Video", height=300)

        with gr.Accordion("Advanced settings", open=False):
            prompt_input = gr.Textbox(
                label="Prompt (optional)",
                value="",
                placeholder="Optional text to guide generation...",
            )
            negative_prompt_input = gr.Textbox(
                label="Negative Prompt (optional)",
                value="",
            )
            seed_input = gr.Number(label="Seed", value=43, precision=0)
            guidance_input = gr.Slider(
                label="Guidance Scale", minimum=1.0, maximum=15.0, value=6.0, step=0.5
            )
            steps_input = gr.Slider(
                label="Inference Steps",
                minimum=10,
                maximum=100,
                value=50,
                step=5,
            )

        gr.Examples(
            examples=[
                ["example1.png", "example1.mp4"],
                ["example2.png", "example2.mp4"],
                ["example3.png", "example3.mp4"],
            ],
            inputs=[input_image, partial_view_video],
            outputs=output_video,
            fn=generate_scene,
            cache_examples=True,
            cache_mode="lazy",
        )

    run_btn.click(
        fn=generate_scene,
        inputs=[
            input_image,
            partial_view_video,
            prompt_input,
            negative_prompt_input,
            seed_input,
            guidance_input,
            steps_input,
        ],
        outputs=output_video,
        api_name="generate",
    )

demo.launch(mcp_server=True)