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)