[DOWNLOAD GIF]

The-Swarm-Corporation · Apr 16, 2024 · d3d8b56 · d3d8b56
1 parent 115c186
commit d3d8b56
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 36 deletions.
diff --git a/servers/text_to_video/test.py b/servers/text_to_video/test.py
@@ -1,4 +1,3 @@
-
 import torch
 from diffusers import (
     AnimateDiffPipeline,
@@ -37,7 +36,6 @@ def text_to_video(
         str: The path to the exported GIF file.
     """
     try:
-
         device = "cuda"
         dtype = torch.float16
 
@@ -46,14 +44,13 @@ def text_to_video(
         base = "emilianJR/epiCRealism"  # Choose to your favorite base model.
         adapter = MotionAdapter().to(device, dtype)
         adapter.load_state_dict(load_file(hf_hub_download(repo, ckpt), device=device))
-        
+
         pipe = AnimateDiffPipeline.from_pretrained(
             base, motion_adapter=adapter, torch_dtype=dtype
         ).to(device)
-        
+
         logger.info(f"Initialized Model: {model_name}")
-
-
+
         pipe.scheduler = EulerDiscreteScheduler.from_config(
             pipe.scheduler.config,
             timestep_spacing="trailing",
@@ -69,12 +66,12 @@ def text_to_video(
         #     )
         #     outputs.append(output)
         #     out = export_to_gif([output], f"{output_path}_{i}.gif")
-            # else:
-            #     out = export_to_video([output], f"{output_path}_{i}.mp4")
+        # else:
+        #     out = export_to_video([output], f"{output_path}_{i}.mp4")
         output = pipe(
-            prompt = task,
-            guidance_scale = guidance_scale,
-            num_inference_steps = inference_steps
+            prompt=task,
+            guidance_scale=guidance_scale,
+            num_inference_steps=inference_steps,
         )
         out = export_to_gif(output.frames[0], output_path)
         return out
@@ -84,4 +81,4 @@ def text_to_video(
 
 
 out = text_to_video(task="A girl in hijab studying in a library")
-print(out)
+print(out)
diff --git a/text_to_video.py b/text_to_video.py
@@ -14,6 +14,7 @@
 from huggingface_hub import hf_hub_download
 from loguru import logger
 from safetensors.torch import load_file
+from fastapi.responses import FileResponse
 
 from swarms_cloud.schema.text_to_video import TextToVideoRequest, TextToVideoResponse
 
@@ -33,6 +34,7 @@
     allow_headers=["*"],
 )
 
+
 def text_to_video(
     task: str,
     model_name: str = "ByteDance/AnimateDiff-Lightning",
@@ -55,7 +57,6 @@ def text_to_video(
         str: The path to the exported GIF file.
     """
     try:
-
         device = "cuda"
         dtype = torch.float16
 
@@ -64,14 +65,13 @@ def text_to_video(
         base = "emilianJR/epiCRealism"  # Choose to your favorite base model.
         adapter = MotionAdapter().to(device, dtype)
         adapter.load_state_dict(load_file(hf_hub_download(repo, ckpt), device=device))
-        
+
         pipe = AnimateDiffPipeline.from_pretrained(
             base, motion_adapter=adapter, torch_dtype=dtype
         ).to(device)
-        
+
         logger.info(f"Initialized Model: {model_name}")
-
-
+
         pipe.scheduler = EulerDiscreteScheduler.from_config(
             pipe.scheduler.config,
             timestep_spacing="trailing",
@@ -87,16 +87,16 @@ def text_to_video(
         #     )
         #     outputs.append(output)
         #     out = export_to_gif([output], f"{output_path}_{i}.gif")
-            # else:
-            #     out = export_to_video([output], f"{output_path}_{i}.mp4")
+        # else:
+        #     out = export_to_video([output], f"{output_path}_{i}.mp4")
         output = pipe(
-            prompt = task,
-            guidance_scale = guidance_scale,
-            num_inference_steps = inference_steps
+            prompt=task,
+            guidance_scale=guidance_scale,
+            num_inference_steps=inference_steps,
         )
-        
+
         logger.info(f"Output ready: {output}")
-        
+
         out = export_to_gif(output.frames[0], output_path)
         logger.info(f"Exported to GIF: {out}")
         return out
@@ -105,7 +105,6 @@ def text_to_video(
         return None
 
 
-
 @app.post("/v1/chat/completions", response_model=TextToVideoResponse)
 async def create_chat_completion(
     request: TextToVideoRequest,  # token: str = Depends(authenticate_user)
@@ -133,20 +132,20 @@ async def create_chat_completion(
         #     logger.error(f"Error: {e}")
         #     raise HTTPException(status_code=500, detail="Internal Server Error")
 
-        out = TextToVideoResponse(
-            status="success",
-            request_details=request,
-            video_url=response,
-            error=None,
-        )
+        # out = TextToVideoResponse(
+        #     status="success",
+        #     request_details=request,
+        #     video_url=response,
+        #     error=None,
+        # )
 
         logger.info(f"Response: {out}")
         logger.info(f"Downloading the file: {response}")
-        # out = FileResponse(
-        #     path=response,
-        #     filename=request.output_path,
-        #     media_type="application/octet-stream",
-        # )
+        out = FileResponse(
+            path=response,
+            filename=request.output_path,
+            media_type="image/gif",  # Use the correct media type for GIFs
+        )
 
         return out
     except Exception as e: