Update app.py

app.py

@@ -1,4 +1,4 @@
-from flask import Flask, request, jsonify, send_file
+import gradio as gr
 import torch
 from diffusers.utils import export_to_video
 from diffusers import AutoencoderKLWan, WanPipeline
@@ -6,13 +6,7 @@ from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepSchedu
 import os
 from uuid import uuid4
 
-app = Flask(__name__)
-
-@app.route('/')
-def index():
-    return jsonify({"message": "Welcome to the Wan2.1 Video Generation API!", "status": "running"})
-
-# Load the model once at startup
+# Load model on startup
 model_id = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
 vae = AutoencoderKLWan.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float32)
 scheduler = UniPCMultistepScheduler(
@@ -25,16 +19,8 @@ pipe = WanPipeline.from_pretrained(model_id, vae=vae, torch_dtype=torch.bfloat16
 pipe.scheduler = scheduler
 pipe.to("cuda")
 
-@app.route('/generate_video', methods=['POST'])
-def generate_video():
-    data = request.json
-    prompt = data.get('prompt')
-    negative_prompt = data.get('negative_prompt', '')
-    height = data.get('height', 720)
-    width = data.get('width', 1280)
-    num_frames = data.get('num_frames', 81)
-    guidance_scale = data.get('guidance_scale', 5.0)
-
+# Define the generation function
+def generate_video(prompt, negative_prompt="", height=720, width=1280, num_frames=81, guidance_scale=5.0):
     output = pipe(
         prompt=prompt,
         negative_prompt=negative_prompt,
@@ -49,7 +35,22 @@ def generate_video():
     os.makedirs("outputs", exist_ok=True)
     export_to_video(output, output_path, fps=16)
 
-    return send_file(output_path, mimetype='video/mp4')
+    return output_path  # Gradio returns this as downloadable file/video
+
+# Gradio Interface
+iface = gr.Interface(
+    fn=generate_video,
+    inputs=[
+        gr.Textbox(label="Prompt"),
+        gr.Textbox(label="Negative Prompt", value=""),
+        gr.Number(label="Height", value=720),
+        gr.Number(label="Width", value=1280),
+        gr.Number(label="Number of Frames", value=81),
+        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."
+)
 
-if __name__ == '__main__':
-    app.run(host='0.0.0.0', port=7860)
+iface.launch()