Create app_t2v.py

app_t2v.py

@@ -0,0 +1,229 @@
+import os
+import sys
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+#import subprocess
+#subprocess.run('pip install flash-attn==2.7.4.post1 --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
+
+# wan2.2-main/gradio_ti2v.py
+import gradio as gr
+import torch
+from huggingface_hub import snapshot_download
+from PIL import Image
+import random
+import numpy as np
+import spaces
+
+import wan
+from wan.configs import WAN_CONFIGS, SIZE_CONFIGS, MAX_AREA_CONFIGS, SUPPORTED_SIZES
+from wan.utils.utils import cache_video
+
+import gc
+
+# --- 1. Global Setup and Model Loading ---
+
+print("Starting Gradio App for Wan 2.2 TI2V-5B...")
+
+# Download model snapshots from Hugging Face Hub
+repo_id = "Wan-AI/Wan2.2-TI2V-5B"
+print(f"Downloading/loading checkpoints for {repo_id}...")
+ckpt_dir = snapshot_download(repo_id, local_dir_use_symlinks=False)
+print(f"Using checkpoints from {ckpt_dir}")
+
+# Load the model configuration
+TASK_NAME = 'ti2v-5B'
+cfg = WAN_CONFIGS[TASK_NAME]
+FIXED_FPS = 24
+MIN_FRAMES_MODEL = 8
+MAX_FRAMES_MODEL = 121 
+
+# Instantiate the pipeline in the global scope
+print("Initializing WanTI2V pipeline...")
+device = "cuda" if torch.cuda.is_available() else "cpu"
+device_id = 0 if torch.cuda.is_available() else -1
+pipeline = wan.WanTI2V(
+    config=cfg,
+    checkpoint_dir=ckpt_dir,
+    device_id=device_id,
+    rank=0,
+    t5_fsdp=False,
+    dit_fsdp=False,
+    use_sp=False,
+    t5_cpu=False,
+    init_on_cpu=False,
+    convert_model_dtype=True,
+)
+print("Pipeline initialized and ready.")
+
+# --- Helper Functions ---
+def select_best_size_for_image(image, available_sizes):
+    """Select the size option with aspect ratio closest to the input image."""
+    if image is None:
+        return available_sizes[0]  # Return first option if no image
+    
+    img_width, img_height = image.size
+    img_aspect_ratio = img_height / img_width
+    
+    best_size = available_sizes[0]
+    best_diff = float('inf')
+    
+    for size_str in available_sizes:
+        # Parse size string like "704*1280" 
+        height, width = map(int, size_str.split('*'))
+        size_aspect_ratio = height / width
+        diff = abs(img_aspect_ratio - size_aspect_ratio)
+        
+        if diff < best_diff:
+            best_diff = diff
+            best_size = size_str
+    
+    return best_size
+
+def handle_image_upload(image):
+    """Handle image upload and return the best matching size."""
+    if image is None:
+        return gr.update()
+    
+    pil_image = Image.fromarray(image).convert("RGB")
+    available_sizes = list(SUPPORTED_SIZES[TASK_NAME])
+    best_size = select_best_size_for_image(pil_image, available_sizes)
+    
+    return gr.update(value=best_size)
+
+def get_duration(
+                 prompt, 
+                 size, 
+                 duration_seconds, 
+                 sampling_steps, 
+                 guide_scale, 
+                 shift, 
+                 seed,
+                 progress):
+    """Calculate dynamic GPU duration based on parameters."""
+    if sampling_steps > 35 and duration_seconds >= 2:
+        return 120
+    elif sampling_steps < 35 or duration_seconds < 2:
+        return 105
+    else:
+        return 90
+
+# --- 2. Gradio Inference Function ---
+@spaces.GPU(duration=get_duration)
+def generate_video(
+    
+    prompt,
+    size,
+    duration_seconds,
+    sampling_steps,
+    guide_scale,
+    shift,
+    seed,
+    progress=gr.Progress(track_tqdm=True)
+):
+    """The main function to generate video, called by the Gradio interface."""
+    if seed == -1:
+        seed = random.randint(0, sys.maxsize)
+
+    # input_image = None
+    # if image is not None:
+    #     input_image = Image.fromarray(image).convert("RGB")
+    #     # Resize image to match selected size
+    #     target_height, target_width = map(int, size.split('*'))
+    #     input_image = input_image.resize((target_width, target_height))
+    
+    # Calculate number of frames based on duration
+    num_frames = np.clip(int(round(duration_seconds * FIXED_FPS)), MIN_FRAMES_MODEL, MAX_FRAMES_MODEL)
+
+    video_tensor = pipeline.generate(
+        input_prompt=prompt,
+        img=None  # Pass None for T2V, Image for I2V
+        size=SIZE_CONFIGS[size],
+        max_area=MAX_AREA_CONFIGS[size],
+        frame_num=num_frames,  # Use calculated frames instead of cfg.frame_num
+        shift=shift,
+        sample_solver='unipc',
+        sampling_steps=int(sampling_steps),
+        guide_scale=guide_scale,
+        seed=seed,
+        offload_model=True
+    )
+
+    # Save the video to a temporary file
+    video_path = cache_video(
+        tensor=video_tensor[None],  # Add a batch dimension
+        save_file=None,  # cache_video will create a temp file
+        fps=cfg.sample_fps,
+        normalize=True,
+        value_range=(-1, 1)
+    )
+    del video_tensor
+    gc.collect()
+    return video_path
+
+
+# --- 3. Gradio Interface ---
+css = ".gradio-container {max-width: 1100px !important; margin: 0 auto} #output_video {height: 500px;} #input_image {height: 500px;}"
+
+with gr.Blocks(css=css, theme=gr.themes.Soft(), delete_cache=(60, 900)) as demo:
+    gr.Markdown("# Wan 2.2 TI2V 5B")
+    gr.Markdown("generate high quality videos using **Wan 2.2 5B Text-Image-to-Video model**,[[model]](https://huggingface.co/Wan-AI/Wan2.2-TI2V-5B),[[paper]](https://arxiv.org/abs/2503.20314)")
+
+    with gr.Row():
+        with gr.Column(scale=2):
+            #image_input = gr.Image(type="numpy", label="Input Image (Optional)", elem_id="input_image")
+            prompt_input = gr.Textbox(label="Prompt", value="A beautiful waterfall in a lush jungle, cinematic.", lines=3)
+            duration_input = gr.Slider(
+                minimum=round(MIN_FRAMES_MODEL/FIXED_FPS, 1), 
+                maximum=round(MAX_FRAMES_MODEL/FIXED_FPS, 1), 
+                step=0.1, 
+                value=2.0, 
+                label="Duration (seconds)", 
+                info=f"Clamped to model's {MIN_FRAMES_MODEL}-{MAX_FRAMES_MODEL} frames at {FIXED_FPS}fps."
+            )
+            size_input = gr.Dropdown(label="Output Resolution", choices=list(SUPPORTED_SIZES[TASK_NAME]), value="704*1280")
+        with gr.Column(scale=2):
+            video_output = gr.Video(label="Generated Video", elem_id="output_video")
+            
+
+            with gr.Accordion("Advanced Settings", open=False):
+                steps_input = gr.Slider(label="Sampling Steps", minimum=10, maximum=50, value=38, step=1)
+                scale_input = gr.Slider(label="Guidance Scale", minimum=1.0, maximum=10.0, value=cfg.sample_guide_scale, step=0.1)
+                shift_input = gr.Slider(label="Sample Shift", minimum=1.0, maximum=20.0, value=cfg.sample_shift, step=0.1)
+                seed_input = gr.Number(label="Seed (-1 for random)", value=-1, precision=0)
+
+            run_button = gr.Button("Generate Video", variant="primary")
+            
+    # Add image upload handler
+    # image_input.upload(
+    #     fn=handle_image_upload,
+    #     inputs=[image_input],
+    #     outputs=[size_input]
+    # )
+    
+    # image_input.clear(
+    #     fn=handle_image_upload,
+    #     inputs=[image_input],
+    #     outputs=[size_input]
+    # )
+
+    # example_image_path = os.path.join(os.path.dirname(__file__), "examples/i2v_input.JPG")
+    # gr.Examples(
+    #     examples=[
+    #         [example_image_path, "The cat removes the glasses from its eyes.", "1280*704", 1.5],
+    #         [None, "A cinematic shot of a boat sailing on a calm sea at sunset.", "1280*704", 2.0],
+    #         [None, "Drone footage flying over a futuristic city with flying cars.", "1280*704", 2.0],
+    #     ],
+        inputs=[prompt_input, size_input, duration_input],
+        outputs=video_output,
+        fn=generate_video,
+        cache_examples=False,
+    )
+
+    run_button.click(
+        fn=generate_video,
+        inputs=[ prompt_input, size_input, duration_input, steps_input, scale_input, shift_input, seed_input],
+        outputs=video_output
+    )
+
+if __name__ == "__main__":
+    demo.launch()