Update app.py

app.py

@@ -1,39 +1,45 @@
 import gradio as gr
 import torch
-from diffusers.utils import export_to_video
-from diffusers import AutoencoderKLWan, WanPipeline
-from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler
-import os
+import ftfy
 from uuid import uuid4
+from diffusers import WanPipeline, AutoencoderKLWan
+from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler
+from diffusers.utils import export_to_video
 
-# Check for available device (CUDA or CPU)
 device = "cuda" if torch.cuda.is_available() else "cpu"
 print(f"Running on {device}...")
 
-# Load the model only once during startup
-try:
-    print("Loading model...")
-    model_id = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
-    vae = AutoencoderKLWan.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float32)
-    scheduler = UniPCMultistepScheduler(
-        prediction_type='flow_prediction',
-        use_flow_sigmas=True,
-        num_train_timesteps=1000,
-        flow_shift=5.0
-    )
-    pipe = WanPipeline.from_pretrained(model_id, vae=vae, torch_dtype=torch.bfloat16)
-    pipe.scheduler = scheduler
-    pipe.to(device)  # Move model to GPU or CPU based on availability
-    print("Model loaded successfully!")
-except Exception as e:
-    print(f"Error loading model: {e}")
-    device = "cpu"  # Fallback to CPU if model loading fails on GPU
-    pipe.to(device)
+# Load model
+model_id = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+vae = AutoencoderKLWan.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float32)
+scheduler = UniPCMultistepScheduler(
+    prediction_type='flow_prediction',
+    use_flow_sigmas=True,
+    num_train_timesteps=1000,
+    flow_shift=5.0
+)
+pipe = WanPipeline.from_pretrained(model_id, vae=vae, torch_dtype=torch.bfloat16)
+pipe.scheduler = scheduler
+pipe.to(device)
 
-# Define the generation function
-def generate_video(prompt, negative_prompt="", height=720, width=1280, num_frames=81, guidance_scale=5.0):
+print("Model loaded successfully.")
+
+def make_divisible_by_16(x):
+    return int(x) - int(x) % 16
+
+def generate_video(prompt, negative_prompt="", height=480, width=832, num_frames=81, guidance_scale=5.0):
     try:
         print(f"Generating video with prompt: {prompt}")
+
+        if not prompt:
+            raise ValueError("Prompt must be provided.")
+
+        # Validate and adjust height/width
+        height = make_divisible_by_16(int(height))
+        width = make_divisible_by_16(int(width))
+        num_frames = int(num_frames)
+        guidance_scale = float(guidance_scale)
+
         output = pipe(
             prompt=prompt,
             negative_prompt=negative_prompt,
@@ -43,32 +49,27 @@ def generate_video(prompt, negative_prompt="", height=720, width=1280, num_frame
             guidance_scale=guidance_scale,
         ).frames[0]
 
-        output_filename = f"{uuid4()}.mp4"
-        output_path = os.path.join("outputs", output_filename)
-        os.makedirs("outputs", exist_ok=True)
+        output_path = f"{uuid4()}.mp4"
         export_to_video(output, output_path, fps=16)
 
-        print(f"Video generated and saved to {output_path}")
-        return output_path  # Gradio returns this as downloadable file/video
+        print(f"Video generated: {output_path}")
+        return output_path
+
     except Exception as e:
         print(f"Error during video generation: {e}")
         return None
 
-# Gradio Interface
 iface = gr.Interface(
     fn=generate_video,
     inputs=[
-        gr.Textbox(label="Prompt"),
+        gr.Textbox(label="Prompt", placeholder="Describe your scene..."),
         gr.Textbox(label="Negative Prompt", value=""),
-        gr.Number(label="Height", value=720),
-        gr.Number(label="Width", value=1280),
+        gr.Number(label="Height", value=480),
+        gr.Number(label="Width", value=832),
         gr.Number(label="Number of Frames", value=81),
-        gr.Number(label="Guidance Scale", value=5.0)
+        gr.Number(label="Guidance Scale", value=5.0),
     ],
     outputs=gr.File(label="Generated Video"),
-    title="Wan2.1 Video Generator",
-    description="Generate realistic videos from text prompts using the Wan2.1 T2V model.",
-    live=True
 )
 
 # Launch Gradio app in API mode