Update app_t2v.py

app_t2v.py

@@ -1,174 +1,69 @@
-# PyTorch 2.8 (temporary hack)
-import os
-os.system('pip install --upgrade --pre --extra-index-url https://download.pytorch.org/whl/nightly/cu126 "torch<2.9" spaces')
-
-# Actual demo code
-import spaces
-import torch
-from diffusers.pipelines.wan.pipeline_wan_i2v import WanImageToVideoPipeline
-from diffusers.models.transformers.transformer_wan import WanTransformer3DModel
-from diffusers.utils.export_utils import export_to_video
 import gradio as gr
+import torch
+from diffusers import WanPipeline, AutoencoderKLWan
+from diffusers.utils import export_to_video
 import tempfile
-import numpy as np
-from PIL import Image
-import random
-
-from optimization import optimize_pipeline_
-
-
-MODEL_ID = "Wan-AI/Wan2.2-T2V-A14B-Diffusers"
-
-LANDSCAPE_WIDTH = 832
-LANDSCAPE_HEIGHT = 480
-MAX_SEED = np.iinfo(np.int32).max
-
-FIXED_FPS = 24
-MIN_FRAMES_MODEL = 8
-MAX_FRAMES_MODEL = 81
-
-MIN_DURATION = round(MIN_FRAMES_MODEL/FIXED_FPS,1)
-MAX_DURATION = round(MAX_FRAMES_MODEL/FIXED_FPS,1)
-
-
-pipe = WanImageToVideoPipeline.from_pretrained(MODEL_ID,
-    transformer=WanTransformer3DModel.from_pretrained('cbensimon/Wan2.2-I2V-A14B-bf16-Diffusers',
-        subfolder='transformer',
-        torch_dtype=torch.bfloat16,
-        device_map='cuda',
-    ),
-    transformer_2=WanTransformer3DModel.from_pretrained('cbensimon/Wan2.2-I2V-A14B-bf16-Diffusers',
-        subfolder='transformer_2',
-        torch_dtype=torch.bfloat16,
-        device_map='cuda',
-    ),
-    torch_dtype=torch.bfloat16,
-).to('cuda')
+import os
 
+# Setup
+dtype = torch.bfloat16
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
-optimize_pipeline_(pipe,
-    image=Image.new('RGB', (LANDSCAPE_WIDTH, LANDSCAPE_HEIGHT)),
-    prompt='prompt',
-    height=LANDSCAPE_HEIGHT,
-    width=LANDSCAPE_WIDTH,
-    num_frames=MAX_FRAMES_MODEL,
+# Load model and VAE once
+vae = AutoencoderKLWan.from_pretrained(
+    "Wan-AI/Wan2.2-T2V-A14B-Diffusers", subfolder="vae", torch_dtype=torch.float32
 )
+pipe = WanPipeline.from_pretrained(
+    "Wan-AI/Wan2.2-T2V-A14B-Diffusers", vae=vae, torch_dtype=dtype
+)
+pipe.to(device)
+
+# Core inference function
+def generate_video(prompt, negative_prompt, height, width, num_frames, guidance_scale, guidance_scale_2, num_steps):
+    with torch.autocast("cuda", dtype=dtype):
+        output = pipe(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            height=height,
+            width=width,
+            num_frames=num_frames,
+            guidance_scale=guidance_scale,
+            guidance_scale_2=guidance_scale_2,
+            num_inference_steps=num_steps,
+        ).frames[0]
+
+    temp_dir = tempfile.mkdtemp()
+    video_path = os.path.join(temp_dir, "output.mp4")
+    export_to_video(output, video_path, fps=16)
+    return video_path
+
+# Gradio UI
+with gr.Blocks() as demo:
+    gr.Markdown("## 🐾 Wan2.2 T2V Demo – Gradio Edition")
 
+    with gr.Row():
+        prompt = gr.Textbox(label="Prompt", value="Two anthropomorphic cats in comfy boxing gear fight intensely.")
+        negative_prompt = gr.Textbox(label="Negative Prompt", value="色调艳丽，过曝，静态，细节模糊不清，字幕，最差质量，丑陋的，多余的手指，畸形")
 
-default_prompt_i2v = "make this image come alive, cinematic motion, smooth animation"
-default_negative_prompt = "色调艳丽, 过曝, 静态, 细节模糊不清, 字幕, 风格, 作品, 画作, 画面, 静止, 整体发灰, 最差质量, 低质量, JPEG压缩残留, 丑陋的, 残缺的, 多余的手指, 画得不好的手部, 画得不好的脸部, 畸形的, 毁容的, 形态畸形的肢体, 手指融合, 静止不动的画面, 杂乱的背景, 三条腿, 背景人很多, 倒着走"
-
-
-def resize_image(image: Image.Image) -> Image.Image:
-    if image.height > image.width:
-        transposed = image.transpose(Image.Transpose.ROTATE_90)
-        resized = resize_image_landscape(transposed)
-        return resized.transpose(Image.Transpose.ROTATE_270)
-    return resize_image_landscape(image)
-
-
-def resize_image_landscape(image: Image.Image) -> Image.Image:
-    target_aspect = LANDSCAPE_WIDTH / LANDSCAPE_HEIGHT
-    width, height = image.size
-    in_aspect = width / height
-    if in_aspect > target_aspect:
-        new_width = round(height * target_aspect)
-        left = (width - new_width) // 2
-        image = image.crop((left, 0, left + new_width, height))
-    else:
-        new_height = round(width / target_aspect)
-        top = (height - new_height) // 2
-        image = image.crop((0, top, width, top + new_height))
-    return image.resize((LANDSCAPE_WIDTH, LANDSCAPE_HEIGHT), Image.LANCZOS)
-
-def get_duration(
-    input_image,
-    prompt,
-    negative_prompt,
-    duration_seconds,
-    guidance_scale,
-    steps,
-    seed,
-    randomize_seed,
-    progress,
-):
-    return steps * 15
-
-@spaces.GPU(duration=get_duration)
-def generate_video(
-    input_image,
-    prompt,
-    negative_prompt=default_negative_prompt,
-    duration_seconds = MAX_DURATION,
-    guidance_scale = 1,
-    steps = 4,
-    seed = 42,
-    randomize_seed = False,
-    progress=gr.Progress(track_tqdm=True),
-):
-    
-    if input_image is None:
-        raise gr.Error("Please upload an input image.")
-    
-    num_frames = np.clip(int(round(duration_seconds * FIXED_FPS)), MIN_FRAMES_MODEL, MAX_FRAMES_MODEL)
-    current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)
-    resized_image = resize_image(input_image)
-
-    output_frames_list = pipe(
-        image=None,
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        height=resized_image.height,
-        width=resized_image.width,
-        num_frames=num_frames,
-        guidance_scale=float(guidance_scale),
-        num_inference_steps=int(steps),
-        generator=torch.Generator(device="cuda").manual_seed(current_seed),
-    ).frames[0]
-
-    with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmpfile:
-        video_path = tmpfile.name
-
-    export_to_video(output_frames_list, video_path, fps=FIXED_FPS)
+    with gr.Row():
+        height = gr.Slider(360, 1024, value=720, step=16, label="Height")
+        width = gr.Slider(360, 1920, value=1280, step=16, label="Width")
 
-    return video_path, current_seed
+    with gr.Row():
+        num_frames = gr.Slider(16, 100, value=81, step=1, label="Number of Frames")
+        num_steps = gr.Slider(10, 60, value=40, step=1, label="Inference Steps")
 
-with gr.Blocks() as demo:
-    gr.Markdown("# Wan2.2-T2V-A14B AND I2V Testing")
-    #gr.Markdown("[CausVid](https://github.com/tianweiy/CausVid) is a distilled version of Wan 2.1 to run faster in just 4-8 steps, [extracted as LoRA by Kijai](https://huggingface.co/Kijai/WanVideo_comfy/blob/main/Wan21_CausVid_14B_T2V_lora_rank32.safetensors) and is compatible with 🧨 diffusers")
     with gr.Row():
-        with gr.Column():
-            input_image_component = gr.Image(type="pil", label="Input Image (auto-resized to target H/W)")
-            prompt_input = gr.Textbox(label="Prompt", value=default_prompt_i2v)
-            duration_seconds_input = gr.Slider(minimum=MIN_DURATION, maximum=MAX_DURATION, step=0.1, value=MAX_DURATION, label="Duration (seconds)", info=f"Clamped to model's {MIN_FRAMES_MODEL}-{MAX_FRAMES_MODEL} frames at {FIXED_FPS}fps.")
-            
-            with gr.Accordion("Advanced Settings", open=False):
-                negative_prompt_input = gr.Textbox(label="Negative Prompt", value=default_negative_prompt, lines=3)
-                seed_input = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42, interactive=True)
-                randomize_seed_checkbox = gr.Checkbox(label="Randomize seed", value=True, interactive=True)
-                steps_slider = gr.Slider(minimum=1, maximum=40, step=1, value=35, label="Inference Steps") 
-                guidance_scale_input = gr.Slider(minimum=0.0, maximum=20.0, step=0.5, value=1.0, label="Guidance Scale", visible=False)
+        guidance_scale = gr.Slider(1.0, 10.0, value=4.0, step=0.5, label="Guidance Scale")
+        guidance_scale_2 = gr.Slider(1.0, 10.0, value=3.0, step=0.5, label="Guidance Scale 2")
 
-            generate_button = gr.Button("Generate Video", variant="primary")
-        with gr.Column():
-            video_output = gr.Video(label="Generated Video", autoplay=True, interactive=False)
-    
-    ui_inputs = [
-        input_image_component, prompt_input,
-        negative_prompt_input, duration_seconds_input,
-        guidance_scale_input, steps_slider, seed_input, randomize_seed_checkbox
-    ]
-    generate_button.click(fn=generate_video, inputs=ui_inputs, outputs=[video_output, seed_input])
+    generate_btn = gr.Button("Generate Video")
+    video_output = gr.Video(label="Generated Video")
 
-    gr.Examples(
-        examples=[ 
-            [
-                "wan_i2v_input.JPG",
-                "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's intricate details and the refreshing atmosphere of the seaside.",
-            ],
-        ],
-        inputs=[input_image_component, prompt_input], outputs=[video_output, seed_input], fn=generate_video, cache_examples="lazy"
+    generate_btn.click(
+        fn=generate_video,
+        inputs=[prompt, negative_prompt, height, width, num_frames, guidance_scale, guidance_scale_2, num_steps],
+        outputs=video_output,
     )
 
-if __name__ == "__main__":
-    demo.queue().launch(mcp_server=True)
+demo.launch()