Update main.py

main.py

@@ -1,8 +1,8 @@
 from fastapi import FastAPI, Query, File, UploadFile
 from fastapi.responses import FileResponse
 import torch
-from diffusion import Diffusion
-from utils import get_id_frame, get_audio_emb, save_video
+from diffusion import Diffusion  # Make sure you import your own modules correctly
+from utils import get_id_frame, get_audio_emb, save_video  # Make sure you import your own modules correctly
 import shutil
 from pathlib import Path
 
@@ -12,18 +12,15 @@ app = FastAPI()
 async def generate_video(
         id_frame_file: UploadFile = File(...),
         audio_file: UploadFile = File(...),
-        gpu: bool = Query(False, description="Use GPU if available"),
+        gpu: bool = Query(True, description="Use GPU if available"),
         id_frame_random: bool = Query(False, description="Pick id_frame randomly from video"),
         inference_steps: int = Query(100, description="Number of inference diffusion steps"),
-        output: str = Query("output.mp4", description="Path to save the output video")
+        output: str = Query("/Users/a/Documents/Automations/git talking heads/output_video.mp4", description="Path to save the output video")
 ):
     device = 'cuda' if gpu and torch.cuda.is_available() else 'cpu'
 
     print('Loading model...')
-    # Load your checkpoint here
-    unet = torch.jit.load("your_checkpoint_path_here")
-    
-    # Replace these arguments with the ones from your original args
+    unet = torch.jit.load("/Users/a/Documents/Automations/git talking heads/checkpoints/crema_script.pt")
     diffusion_args = {
         "in_channels": 3,
         "image_size": 128,
@@ -43,9 +40,15 @@ async def generate_video(
         shutil.copyfileobj(audio_file.file, buffer)
 
     id_frame = get_id_frame(str(id_frame_path), random=id_frame_random, resize=diffusion_args["image_size"]).to(device)
-    audio, audio_emb = get_audio_emb(str(audio_path), "your_encoder_checkpoint_here", device)
+    audio, audio_emb = get_audio_emb(str(audio_path), "/Users/a/Documents/Automations/git talking heads/checkpoints/audio_encoder.pt", device)
 
-    samples = diffusion.sample(id_frame, audio_emb.unsqueeze(0))
+    unet_args = {
+        "n_audio_motion_embs": 2,
+        "n_motion_frames": 2,
+        "motion_channels": 3
+    }
+    
+    samples = diffusion.sample(id_frame, audio_emb.unsqueeze(0), **unet_args)
 
     save_video(output, samples, audio=audio, fps=25, audio_rate=16000)
     print(f'Results saved at {output}')