Update app.py

app.py

@@ -1,10 +1,15 @@
-from fastapi import FastAPI, UploadFile, File, Form
+from fastapi import FastAPI, UploadFile, File, Form, BackgroundTasks
 import os, uuid, subprocess, torch, cv2
 import whisper
 from scenedetect import VideoManager, SceneManager
 from scenedetect.detectors import ContentDetector
 from ultralytics import YOLO
 from diffusers import StableVideoDiffusionPipeline
+from PIL import Image
+
+# ===============================
+# App + dirs
+# ===============================
 
 app = FastAPI()
 
@@ -13,20 +18,35 @@ OUTPUT_DIR = "outputs"
 os.makedirs(UPLOAD_DIR, exist_ok=True)
 os.makedirs(OUTPUT_DIR, exist_ok=True)
 
+# ===============================
+# Device / dtype (CRITICAL)
+# ===============================
+
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32
 
-# ===== Load models =====
+# ===============================
+# Job store
+# ===============================
 
-whisper_model = whisper.load_model("base")
+jobs = {}
 
+# ===============================
+# Load models (safe on startup)
+# ===============================
+
+whisper_model = whisper.load_model("base")
 yolo = YOLO("yolov8n.pt")
 
 svd = StableVideoDiffusionPipeline.from_pretrained(
     "stabilityai/stable-video-diffusion-img2vid",
-    torch_dtype=torch.float16
-).to(DEVICE)
+    dtype=DTYPE
+)
+svd.to(DEVICE)
 
-# ===== Endpoints =====
+# ===============================
+# Endpoints
+# ===============================
 
 @app.post("/captions")
 async def captions(file: UploadFile = File(...)):
@@ -87,46 +107,72 @@ async def smart_crop(
         "aspect": aspect
     }
 
+# ===============================
+# Background job
+# ===============================
+
+def run_edit_job(job_id: str, video_path: str, frame_path: str):
+    try:
+        subprocess.run(
+            [
+                "ffmpeg", "-y",
+                "-i", video_path,
+                "-vf", "scale=512:512:force_original_aspect_ratio=decrease",
+                "-frames:v", "1",
+                frame_path
+            ],
+            check=True
+        )
+
+        img = Image.open(frame_path).convert("RGB")
+
+        with torch.no_grad():
+            output = svd(
+                image=img,
+                num_frames=16,
+                decode_chunk_size=8
+            )
+
+        jobs[job_id]["status"] = "done"
+        jobs[job_id]["frames"] = len(output.frames)
+
+    except Exception as e:
+        jobs[job_id]["status"] = "error"
+        jobs[job_id]["error"] = str(e)
+
+
+@app.get("/status/{job_id}")
+def job_status(job_id: str):
+    return jobs.get(job_id, {"status": "not_found"})
+
+
 @app.post("/edit")
 async def edit_video(
+    background_tasks: BackgroundTasks,
     file: UploadFile = File(...),
-    prompt: str = Form(...)  # keep prompt for future use
+    prompt: str = Form(...)
 ):
-    uid = uuid.uuid4().hex
-    video_path = os.path.join(UPLOAD_DIR, f"{uid}.mp4")
-    frame_path = os.path.join(OUTPUT_DIR, f"{uid}.png")
+    job_id = uuid.uuid4().hex
+
+    video_path = os.path.join(UPLOAD_DIR, f"{job_id}.mp4")
+    frame_path = os.path.join(OUTPUT_DIR, f"{job_id}.png")
 
-    # Save video
     with open(video_path, "wb") as f:
         f.write(await file.read())
 
-    # Extract first frame safely
-    subprocess.run(
-        [
-            "ffmpeg",
-            "-y",
-            "-i", video_path,
-            "-vf", "scale=512:512:force_original_aspect_ratio=decrease",
-            "-frames:v", "1",
-            frame_path
-        ],
-        check=True
-    )
-
-    from PIL import Image
-    img = Image.open(frame_path).convert("RGB")
-
-    # IMPORTANT: SVD does NOT take prompt
-    with torch.no_grad():
-        output = svd(
-            image=img,
-            num_frames=16,
-            decode_chunk_size=8
-        )
+    jobs[job_id] = {
+        "status": "running",
+        "prompt_received_but_unused": prompt
+    }
 
-    frames = output.frames
+    background_tasks.add_task(
+        run_edit_job,
+        job_id,
+        video_path,
+        frame_path
+    )
 
     return {
-        "prompt_received_but_unused": prompt,
-        "frames_generated": len(frames)
+        "job_id": job_id,
+        "status": "running"
     }