Update app.py

app.py

@@ -1,23 +1,20 @@
 import torch
-from diffusers import AutoencoderKLWan, WanImageToVideoPipeline, UniPCMultistepScheduler
+from diffusers import AutoencoderKLWan, WanTextToVideoPipeline, UniPCMultistepScheduler
 from diffusers.utils import export_to_video
-from transformers import CLIPVisionModel
 import gradio as gr
 import tempfile
 import spaces
 from huggingface_hub import hf_hub_download
 import numpy as np
-from PIL import Image
 import random
 
-MODEL_ID = "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"
+MODEL_ID = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
 LORA_REPO_ID = "Kijai/WanVideo_comfy"
-LORA_FILENAME = "Wan21_CausVid_14B_T2V_lora_rank32.safetensors"
+LORA_FILENAME = "Wan21_CausVid_bidirect2_T2V_1_3B_lora_rank32.safetensors"
 
-image_encoder = CLIPVisionModel.from_pretrained(MODEL_ID, subfolder="image_encoder", torch_dtype=torch.float32)
 vae = AutoencoderKLWan.from_pretrained(MODEL_ID, subfolder="vae", torch_dtype=torch.float32)
-pipe = WanImageToVideoPipeline.from_pretrained(
-    MODEL_ID, vae=vae, image_encoder=image_encoder, torch_dtype=torch.bfloat16
+pipe = WanTextToVideoPipeline.from_pretrained(
+    MODEL_ID, vae=vae, torch_dtype=torch.bfloat16
 )
 pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=8.0)
 pipe.to("cuda")
@@ -30,7 +27,6 @@ pipe.fuse_lora()
 MOD_VALUE = 32
 DEFAULT_H_SLIDER_VALUE = 512
 DEFAULT_W_SLIDER_VALUE = 896
-NEW_FORMULA_MAX_AREA = 480.0 * 832.0 
 
 SLIDER_MIN_H, SLIDER_MAX_H = 128, 896
 SLIDER_MIN_W, SLIDER_MAX_W = 128, 896
@@ -40,46 +36,10 @@ FIXED_FPS = 24
 MIN_FRAMES_MODEL = 8
 MAX_FRAMES_MODEL = 81 
 
-default_prompt_i2v = "make this image come alive, cinematic motion, smooth animation"
+default_prompt_t2v = "a beautiful sunset over mountains, cinematic, smooth camera movement"
 default_negative_prompt = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards, watermark, text, signature"
 
-
-def _calculate_new_dimensions_wan(pil_image, mod_val, calculation_max_area,
-                                 min_slider_h, max_slider_h,
-                                 min_slider_w, max_slider_w,
-                                 default_h, default_w):
-    orig_w, orig_h = pil_image.size
-    if orig_w <= 0 or orig_h <= 0:
-        return default_h, default_w
-
-    aspect_ratio = orig_h / orig_w
-    
-    calc_h = round(np.sqrt(calculation_max_area * aspect_ratio))
-    calc_w = round(np.sqrt(calculation_max_area / aspect_ratio))
-
-    calc_h = max(mod_val, (calc_h // mod_val) * mod_val)
-    calc_w = max(mod_val, (calc_w // mod_val) * mod_val)
-    
-    new_h = int(np.clip(calc_h, min_slider_h, (max_slider_h // mod_val) * mod_val))
-    new_w = int(np.clip(calc_w, min_slider_w, (max_slider_w // mod_val) * mod_val))
-    
-    return new_h, new_w
-
-def handle_image_upload_for_dims_wan(uploaded_pil_image, current_h_val, current_w_val):
-    if uploaded_pil_image is None:
-        return gr.update(value=DEFAULT_H_SLIDER_VALUE), gr.update(value=DEFAULT_W_SLIDER_VALUE)
-    try:
-        new_h, new_w = _calculate_new_dimensions_wan(
-            uploaded_pil_image, MOD_VALUE, NEW_FORMULA_MAX_AREA,
-            SLIDER_MIN_H, SLIDER_MAX_H, SLIDER_MIN_W, SLIDER_MAX_W,
-            DEFAULT_H_SLIDER_VALUE, DEFAULT_W_SLIDER_VALUE
-        )
-        return gr.update(value=new_h), gr.update(value=new_w)
-    except Exception as e:
-        gr.Warning("Error attempting to calculate new dimensions")
-        return gr.update(value=DEFAULT_H_SLIDER_VALUE), gr.update(value=DEFAULT_W_SLIDER_VALUE)
-
-def get_duration(input_image, prompt, height, width, 
+def get_duration(prompt, height, width, 
                    negative_prompt, duration_seconds,
                    guidance_scale, steps,
                    seed, randomize_seed, 
@@ -92,21 +52,20 @@ def get_duration(input_image, prompt, height, width,
         return 60
 
 @spaces.GPU(duration=get_duration)
-def generate_video(input_image, prompt, height, width, 
+def generate_video(prompt, height, width, 
                    negative_prompt=default_negative_prompt, duration_seconds = 2,
                    guidance_scale = 1, steps = 4,
                    seed = 42, randomize_seed = False, 
                    progress=gr.Progress(track_tqdm=True)):
     """
-    Generate a video from an input image using the Wan 2.1 I2V model with CausVid LoRA.
+    Generate a video from a text prompt using the Wan 2.1 T2V model with CausVid LoRA.
     
-    This function takes an input image and generates a video animation based on the provided
-    prompt and parameters. It uses the Wan 2.1 14B Image-to-Video model with CausVid LoRA
+    This function takes a text prompt and generates a video based on the provided
+    prompt and parameters. It uses the Wan 2.1 1.3B Text-to-Video model with CausVid LoRA
     for fast generation in 4-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.
+        prompt (str): Text prompt describing the desired video content.
         height (int): Target height for the output video. Will be adjusted to multiple of MOD_VALUE (32).
         width (int): Target width for the output video. Will be adjusted to multiple of MOD_VALUE (32).
         negative_prompt (str, optional): Negative prompt to avoid unwanted elements. 
@@ -129,17 +88,16 @@ def generate_video(input_image, prompt, height, width,
             - current_seed (int): The seed used for generation (useful when randomize_seed=True)
     
     Raises:
-        gr.Error: If input_image is None (no image uploaded).
+        gr.Error: If prompt is empty or None.
     
     Note:
-        - The function automatically resizes the input image to the target dimensions
         - Frame count is calculated as duration_seconds * FIXED_FPS (24)
         - Output dimensions are adjusted to be multiples of MOD_VALUE (32)
         - 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.")
+    if not prompt or prompt.strip() == "":
+        raise gr.Error("Please enter a text prompt.")
 
     target_h = max(MOD_VALUE, (int(height) // MOD_VALUE) * MOD_VALUE)
     target_w = max(MOD_VALUE, (int(width) // MOD_VALUE) * MOD_VALUE)
@@ -148,11 +106,9 @@ def generate_video(input_image, prompt, height, width,
     
     current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)
 
-    resized_image = input_image.resize((target_w, target_h))
-
     with torch.inference_mode():
         output_frames_list = pipe(
-            image=resized_image, prompt=prompt, negative_prompt=negative_prompt,
+            prompt=prompt, negative_prompt=negative_prompt,
             height=target_h, width=target_w, num_frames=num_frames,
             guidance_scale=float(guidance_scale), num_inference_steps=int(steps),
             generator=torch.Generator(device="cuda").manual_seed(current_seed)
@@ -164,12 +120,11 @@ def generate_video(input_image, prompt, height, width,
     return video_path, current_seed
 
 with gr.Blocks() as demo:
-    gr.Markdown("# Fast 4 steps Wan 2.1 I2V (14B) with CausVid LoRA")
-    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")
+    gr.Markdown("# Fast 4 steps Wan 2.1 T2V (1.3B) with CausVid LoRA")
+    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_bidirect2_T2V_1_3B_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)
+            prompt_input = gr.Textbox(label="Prompt", value=default_prompt_t2v, placeholder="Describe the video you want to generate...")
             duration_seconds_input = gr.Slider(minimum=round(MIN_FRAMES_MODEL/FIXED_FPS,1), maximum=round(MAX_FRAMES_MODEL/FIXED_FPS,1), step=0.1, value=2, 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):
@@ -185,21 +140,9 @@ with gr.Blocks() as demo:
             generate_button = gr.Button("Generate Video", variant="primary")
         with gr.Column():
             video_output = gr.Video(label="Generated Video", autoplay=True, interactive=False)
-
-    input_image_component.upload(
-        fn=handle_image_upload_for_dims_wan,
-        inputs=[input_image_component, height_input, width_input],
-        outputs=[height_input, width_input]
-    )
-    
-    input_image_component.clear( 
-        fn=handle_image_upload_for_dims_wan,
-        inputs=[input_image_component, height_input, width_input],
-        outputs=[height_input, width_input]
-    )
     
     ui_inputs = [
-        input_image_component, prompt_input, height_input, width_input,
+        prompt_input, height_input, width_input,
         negative_prompt_input, duration_seconds_input,
         guidance_scale_input, steps_slider, seed_input, randomize_seed_checkbox
     ]
@@ -207,10 +150,11 @@ with gr.Blocks() as demo:
 
     gr.Examples(
         examples=[ 
-            ["peng.png", "a penguin playfully dancing in the snow, Antarctica", 896, 512],
-            ["forg.jpg", "the frog jumps around", 448, 832],
+            ["a majestic eagle soaring through mountain peaks, cinematic aerial view", 896, 512],
+            ["a serene ocean wave crashing on a sandy beach at sunset", 448, 832],
+            ["a field of flowers swaying in the wind, spring morning light", 512, 896],
         ],
-        inputs=[input_image_component, prompt_input, height_input, width_input], outputs=[video_output, seed_input], fn=generate_video, cache_examples="lazy"
+        inputs=[prompt_input, height_input, width_input], outputs=[video_output, seed_input], fn=generate_video, cache_examples="lazy"
     )
 
 if __name__ == "__main__":