Update app.py

app.py

@@ -5,7 +5,7 @@ os.system('pip install --upgrade --pre --extra-index-url https://download.pytorc
 # Actual demo code
 import spaces
 import torch
-from diffusers.pipelines.wan.pipeline_wan import WanPipeline
+from diffusers import WanPipeline, AutoencoderKLWan
 from diffusers.models.transformers.transformer_wan import WanTransformer3DModel
 from diffusers.utils.export_utils import export_to_video
 import gradio as gr
@@ -30,7 +30,7 @@ MAX_FRAMES_MODEL = 81
 MIN_DURATION = round(MIN_FRAMES_MODEL/FIXED_FPS,1)
 MAX_DURATION = round(MAX_FRAMES_MODEL/FIXED_FPS,1)
 
-
+vae = AutoencoderKLWan.from_pretrained("Wan-AI/Wan2.2-T2V-A14B-Diffusers", subfolder="vae", torch_dtype=torch.float32)
 pipe = WanPipeline.from_pretrained(MODEL_ID,
     transformer=WanTransformer3DModel.from_pretrained('cbensimon/Wan2.2-I2V-A14B-bf16-Diffusers',
         subfolder='transformer',
@@ -42,6 +42,7 @@ pipe = WanPipeline.from_pretrained(MODEL_ID,
         torch_dtype=torch.bfloat16,
         device_map='cuda',
     ),
+    vae=vae,
     torch_dtype=torch.bfloat16,
 ).to('cuda')
 
@@ -80,7 +81,6 @@ for i in range(3):
     torch.cuda.empty_cache()
 
 optimize_pipeline_(pipe,
-    # image=Image.new('RGB', (LANDSCAPE_WIDTH, LANDSCAPE_HEIGHT)),
     prompt='prompt',
     height=LANDSCAPE_HEIGHT,
     width=LANDSCAPE_WIDTH,
@@ -88,34 +88,11 @@ optimize_pipeline_(pipe,
 )
 
 
-default_prompt_i2v = "make this image come alive, cinematic motion, smooth animation"
+default_prompt_i2v = "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage."
 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,
@@ -130,13 +107,12 @@ def get_duration(
 
 @spaces.GPU(duration=get_duration)
 def generate_video(
-    input_image,
     prompt,
     negative_prompt=default_negative_prompt,
     duration_seconds = MAX_DURATION,
     guidance_scale = 1,
     guidance_scale_2 = 3,
-    steps = 6,
+    steps = 4,
     seed = 42,
     randomize_seed = False,
     progress=gr.Progress(track_tqdm=True),
@@ -149,7 +125,6 @@ def generate_video(
     for fast generation in 6-8 steps.
     
     Args:
-        input_image (PIL.Image): The input image to animate. Will be resized to target dimensions.
         prompt (str): Text prompt describing the desired animation or motion.
         negative_prompt (str, optional): Negative prompt to avoid unwanted elements. 
             Defaults to default_negative_prompt (contains unwanted visual artifacts).
@@ -182,15 +157,11 @@ def generate_video(
         - The function uses GPU acceleration via the @spaces.GPU decorator
         - Generation time varies based on steps and duration (see get_duration function)
     """
-    # 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=resized_image,
         prompt=prompt,
         negative_prompt=negative_prompt,
         height=480,
@@ -214,7 +185,6 @@ with gr.Blocks() as demo:
     gr.Markdown("run Wan 2.2 in just 6-8 steps, with [FusionX Phantom LoRA by DeeJayT](https://huggingface.co/vrgamedevgirl84/Wan14BT2VFusioniX/tree/main/FusionX_LoRa), 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)", visible=False)
             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.")
             
@@ -237,15 +207,14 @@ with gr.Blocks() as demo:
     ]
     generate_button.click(fn=generate_video, inputs=ui_inputs, outputs=[video_output, seed_input])
 
-    # 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"
-    # )
+    gr.Examples(
+        examples=[ 
+            [
+                "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=[prompt_input], outputs=[video_output, seed_input], fn=generate_video, cache_examples="lazy"
+    )
 
 if __name__ == "__main__":
     demo.queue().launch(mcp_server=True)