Update main.py

main.py

@@ -1,13 +1,14 @@
 import torch
-import os
 import uuid
-from diffusers import AnimateDiffPipeline, EulerDiscreteScheduler
+
+from diffusers import AnimateDiffPipeline, MotionAdapter, EulerDiscreteScheduler
 from diffusers.utils import export_to_video
 from huggingface_hub import hf_hub_download
 from safetensors.torch import load_file
-from transformers import CLIPFeatureExtractor
-from fastapi import FastAPI, Form, HTTPException
-from fastapi.responses import FileResponse
+from PIL import Image
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+import uvicorn
 
 app = FastAPI()
 
@@ -18,16 +19,15 @@ bases = {
     "3d": "Lykon/DreamShaper",
     "Anime": "Yntec/mistoonAnime2"
 }
-motion_choices = {
-    "": "Default",
-    "guoyww/animatediff-motion-lora-zoom-in": "Zoom in",
-    "guoyww/animatediff-motion-lora-zoom-out": "Zoom out",
-    "guoyww/animatediff-motion-lora-tilt-up": "Tilt up",
-    "guoyww/animatediff-motion-lora-tilt-down": "Tilt down",
-    "guoyww/animatediff-motion-lora-pan-left": "Pan left",
-    "guoyww/animatediff-motion-lora-pan-right": "Pan right",
-    "guoyww/animatediff-motion-lora-rolling-anticlockwise": "Roll left",
-    "guoyww/animatediff-motion-lora-rolling-clockwise": "Roll right"
+motions = {
+    "Zoom in": "guoyww/animatediff-motion-lora-zoom-in",
+    "Zoom out": "guoyww/animatediff-motion-lora-zoom-out",
+    "Tilt up": "guoyww/animatediff-motion-lora-tilt-up",
+    "Tilt down": "guoyww/animatediff-motion-lora-tilt-down",
+    "Pan left": "guoyww/animatediff-motion-lora-pan-left",
+    "Pan right": "guoyww/animatediff-motion-lora-pan-right",
+    "Roll left": "guoyww/animatediff-motion-lora-rolling-anticlockwise",
+    "Roll right": "guoyww/animatediff-motion-lora-rolling-clockwise",
 }
 step_loaded = None
 base_loaded = "Realistic"
@@ -43,14 +43,28 @@ pipe = AnimateDiffPipeline.from_pretrained(bases[base_loaded], torch_dtype=dtype
 pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing", beta_schedule="linear")
 
 # Safety checkers
+from transformers import CLIPFeatureExtractor
+
 feature_extractor = CLIPFeatureExtractor.from_pretrained("openai/clip-vit-base-patch32")
 
-# Function to generate image
-def generate_image(prompt, base="Realistic", motion="", step=8):
+class GenerateImageRequest(BaseModel):
+    prompt: str
+    base: str = "Realistic"
+    motion: str = ""
+    step: int = 8
+
+@app.post("/generate-image")
+def generate_image(request: GenerateImageRequest):
     global step_loaded
     global base_loaded
     global motion_loaded
-    print(prompt, base, step, motion)
+
+    prompt = request.prompt
+    base = request.base
+    motion = request.motion
+    step = request.step
+
+    print(prompt, base, step)
 
     if step_loaded != step:
         repo = "ByteDance/AnimateDiff-Lightning"
@@ -64,8 +78,9 @@ def generate_image(prompt, base="Realistic", motion="", step=8):
 
     if motion_loaded != motion:
         pipe.unload_lora_weights()
-        if motion != "":
-            pipe.load_lora_weights(motion, adapter_name="motion")
+        if motion in motions:
+            motion_repo = motions[motion]
+            pipe.load_lora_weights(motion_repo, adapter_name="motion")
             pipe.set_adapters(["motion"], [0.7])
         motion_loaded = motion
 
@@ -74,23 +89,8 @@ def generate_image(prompt, base="Realistic", motion="", step=8):
     name = str(uuid.uuid4()).replace("-", "")
     path = f"/tmp/{name}.mp4"
     export_to_video(output.frames[0], path, fps=10)
-    return path
-
-# API Endpoint to generate video
-@app.post("/generate-video/")
-async def generate_video(
-    prompt: str = Form(...), 
-    base: str = Form("Realistic"), 
-    motion: str = Form(""), 
-    step: int = Form(8)
-):
-    try:
-        video_path = generate_image(prompt, base, motion, step)
-        return FileResponse(video_path, media_type="video/mp4")
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
-
-# Run the app
+    
+    return {"video_path": path}
+
 if __name__ == "__main__":
-    import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=7860)