Update app.py

app.py

@@ -1,41 +1,37 @@
-import gradio as gr
+from fastapi import FastAPI, Form
+from fastapi.responses import FileResponse
 import torch
 from diffusers import StableDiffusionPipeline, StableVideoDiffusionPipeline
 from diffusers.utils import export_to_video
 from PIL import Image
 
-# Load models
+app = FastAPI()
+
+# Check device
+device = "cuda" if torch.cuda.is_available() else "cpu"
+dtype = torch.float16 if device == "cuda" else torch.float32
+variant = "fp16" if device == "cuda" else None
+
+# Load text-to-image model
 text2img_pipe = StableDiffusionPipeline.from_pretrained(
-    "runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16
-).to("cuda")
-text2img_pipe.enable_model_cpu_offload()
+    "runwayml/stable-diffusion-v1-5",
+    torch_dtype=dtype
+).to(device)
 
+# Load image-to-video model
 video_pipe = StableVideoDiffusionPipeline.from_pretrained(
-    "stabilityai/stable-video-diffusion-img2vid-xt", torch_dtype=torch.float16, variant="fp16"
-).to("cuda")
-video_pipe.enable_model_cpu_offload()
+    "stabilityai/stable-video-diffusion-img2vid-xt",
+    torch_dtype=dtype,
+    variant=variant
+).to(device)
 
-def generate_video(prompt, aspect_ratio):
-    # Generate image
+@app.post("/generate")
+def generate_video(prompt: str = Form(...), aspect: str = Form("16:9")):
     image = text2img_pipe(prompt).images[0]
-    if aspect_ratio == "16:9":
-        image = image.resize((1024, 576))
-    else:
-        image = image.resize((576, 1024))
+    image = image.resize((1024, 576) if aspect == "16:9" else (576, 1024))
 
-    # Generate video
     generator = torch.manual_seed(42)
     frames = video_pipe(image, decode_chunk_size=8, generator=generator).frames[0]
     export_to_video(frames, "output.mp4", fps=7)
-    return "output.mp4"
 
-gr.Interface(
-    fn=generate_video,
-    inputs=[
-        gr.Textbox(label="Scene Prompt", placeholder="e.g. A 3D Islamic city at sunset, 4K"),
-        gr.Radio(["16:9", "9:16"], label="Aspect Ratio", value="16:9")
-    ],
-    outputs=gr.Video(label="Generated Video"),
-    title="Text → Image → Video Generator",
-    description="Enter a scene prompt and get a cinematic video. Powered by Stable Diffusion + Stable Video Diffusion."
-).launch()
+    return FileResponse("output.mp4", media_type="video/mp4", filename="output.mp4")