Update main.py

main.py

@@ -7,46 +7,39 @@ from PIL import Image
 import ffmpeg
 import tempfile
 import os
+from huggingface_hub import InferenceClient
 
 app = FastAPI()
 
-API_URL = "https://api-inference.huggingface.co/models/runwayml/stable-diffusion-v1-5"
+client = InferenceClient("stabilityai/stable-video-diffusion-img2vid-xt-1-1-tensorrt")
 
 @app.post("/generate_video/")
 async def generate_video_api(
-    file: UploadFile = File(...),
-    prompt: str = Form(...),
-    num_frames: int = Form(14),
-    fps: int = Form(7)
+        file: UploadFile = File(...),
+        num_frames: int = Form(14),
+        fps: int = Form(7)
 ):
     try:
         # Read the uploaded image file
         image_content = await file.read()
-        
-        # Generate a sequence of images
-        frames = []
-        for i in range(num_frames):
-            # Prepare the payload
-            payload = {
-                "inputs": f"{prompt}, frame {i+1} of {num_frames}",
+        image = Image.open(io.BytesIO(image_content))
+
+        # Generate video frames using the stable-video-diffusion model
+        video_frames = client.post(
+            json={
+                "inputs": image,
                 "parameters": {
-                    "num_inference_steps": 50,
-                    "guidance_scale": 7.5,
+                    "num_inference_steps": 25,
+                    "num_frames": num_frames,
                 }
             }
-
-            # Make the API call
-            response = requests.post(API_URL, json=payload)
-            response.raise_for_status()
-
-            # Add the generated image to frames
-            frames.append(Image.open(io.BytesIO(response.content)))
+        )
 
         # Create a temporary directory
         with tempfile.TemporaryDirectory() as tmpdir:
             # Save frames as temporary files
             frame_files = []
-            for i, frame in enumerate(frames):
+            for i, frame in enumerate(video_frames):
                 frame_file = os.path.join(tmpdir, f"frame_{i:03d}.png")
                 frame.save(frame_file)
                 frame_files.append(frame_file)
@@ -70,17 +63,8 @@ async def generate_video_api(
         # Return the video as a streaming response
         return StreamingResponse(io.BytesIO(video_content), media_type="video/mp4")
 
-    except requests.exceptions.HTTPError as http_err:
-        # Handle HTTP errors
-        error_detail = f"HTTP error occurred: {http_err}"
-        raise HTTPException(status_code=500, detail=error_detail)
-
-    except ffmpeg.Error as e:
-        # Handle FFmpeg errors
-        raise HTTPException(status_code=500, detail=f"FFmpeg error occurred: {e.stderr.decode()}")
-
     except Exception as err:
-        # Handle any other errors
+        # Handle any errors
         raise HTTPException(status_code=500, detail=f"An error occurred: {err}")
 
 if __name__ == "__main__":