Update main.py

main.py

@@ -1,44 +1,42 @@
-from fastapi import FastAPI
+from fastapi import FastAPI, File, UploadFile
 from pydantic import BaseModel
 from huggingface_hub import InferenceClient
 import uvicorn
 from fastapi.responses import StreamingResponse
 import io
+import base64
 
 app = FastAPI()
 
-client = InferenceClient("stabilityai/stable-diffusion-2-1")
+client = InferenceClient("stabilityai/stable-video-diffusion-img2vid-xt-1-1-tensorrt")
 
 class Item(BaseModel):
-    prompt: str
-    negative_prompt: str = ""
-    num_inference_steps: int = 50
-    guidance_scale: float = 7.5
-    width: int = 512
-    height: int = 512
+    fps: int = 7
+    num_frames: int = 14
+    motion_bucket_id: int = 127
+    cond_aug: float = 0.02
+    seed: int = 0
 
-def generate_image(item: Item):
-    image = client.text_to_image(
-        prompt=item.prompt,
-        negative_prompt=item.negative_prompt,
-        num_inference_steps=item.num_inference_steps,
-        guidance_scale=item.guidance_scale,
-        width=item.width,
-        height=item.height,
-    )
-    return image
-
-@app.post("/generate_image/")
-async def generate_image_api(item: Item):
-    image = generate_image(item)
+@app.post("/generate_video/")
+async def generate_video_api(item: Item, file: UploadFile = File(...)):
+    # Read the uploaded image file
+    image_content = await file.read()
     
-    # Convert the image to a byte stream
-    img_byte_arr = io.BytesIO()
-    image.save(img_byte_arr, format='PNG')
-    img_byte_arr.seek(0)
+    # Encode the image to base64
+    image_base64 = base64.b64encode(image_content).decode('utf-8')
+
+    # Generate the video
+    video = client.image_to_video(
+        image=image_base64,
+        fps=item.fps,
+        num_frames=item.num_frames,
+        motion_bucket_id=item.motion_bucket_id,
+        cond_aug=item.cond_aug,
+        seed=item.seed
+    )
     
-    # Return the image as a streaming response
-    return StreamingResponse(img_byte_arr, media_type="image/png")
+    # Return the video as a streaming response
+    return StreamingResponse(io.BytesIO(video), media_type="video/mp4")
 
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)