Update app.py

app.py

@@ -1,17 +1,18 @@
 # app.py
-from fastapi import FastAPI, Query
+from fastapi import FastAPI
 from fastapi.staticfiles import StaticFiles
 from pydantic import BaseModel
-from huggingface_hub import hf_hub_download
 import uvicorn
 import os
 import uuid
-import shutil
-import requests
+
+# GGUF & video libraries
+from gguf_runtime import GGUFModel  # Make sure gguf-runtime is installed
+from moviepy.editor import ImageSequenceClip
 
 app = FastAPI(title="WAN2 GGUF API", version="1.0")
 
-# ✅ Directories
+# -------------------- Directories --------------------
 MODEL_REPO = "calcuis/wan2-gguf"
 MODEL_FILE = "wan2.2-animate-14b-q4_0.gguf"
 MODEL_DIR = "models"
@@ -20,7 +21,9 @@ OUTPUT_DIR = "outputs"
 os.makedirs(MODEL_DIR, exist_ok=True)
 os.makedirs(OUTPUT_DIR, exist_ok=True)
 
-# ✅ Download model file at startup
+# -------------------- Download model --------------------
+from huggingface_hub import hf_hub_download
+
 model_path = hf_hub_download(
     repo_id=MODEL_REPO,
     filename=MODEL_FILE,
@@ -28,70 +31,59 @@ model_path = hf_hub_download(
 )
 print("✅ Model downloaded to:", model_path)
 
+# -------------------- Load GGUF model --------------------
+print("Loading WAN2 GGUF model...")
+wan2_model = GGUFModel(model_path)
+print("✅ WAN2 model loaded")
 
-# Request schema
+# -------------------- Request schema --------------------
 class PromptRequest(BaseModel):
     prompt: str
     steps: int = 20
 
-
+# -------------------- Routes --------------------
 @app.get("/")
 def root():
     return {"message": "WAN2 GGUF API is running!"}
 
+@app.get("/generate")
+def generate_video_get(q: str, steps: int = 20):
+    """Allows GET requests with ?q=... for browser testing"""
+    return generate_video(PromptRequest(prompt=q, steps=steps))
 
-def generate_output(prompt: str, steps: int = 20):
-    """Helper that creates a dummy video and returns JSON"""
+@app.post("/generate")
+def generate_video(request: PromptRequest):
+    """Generates video from prompt using WAN2 GGUF"""
     # Unique filename
     file_id = str(uuid.uuid4())
     file_path = os.path.join(OUTPUT_DIR, f"{file_id}.mp4")
 
-    # Hugging Face placeholder video
-    placeholder_url = (
-        "https://huggingface.co/datasets/huggingface/documentation-images/"
-        "resolve/main/video-placeholder.mp4"
+    # ---- WAN2 inference ----
+    # This generates frames from prompt
+    frames = wan2_model.generate_video(
+        prompt=request.prompt,
+        steps=request.steps
     )
 
-    # Download placeholder if not already present
-    placeholder_file = os.path.join(OUTPUT_DIR, "placeholder.mp4")
-    if not os.path.exists(placeholder_file):
-        r = requests.get(placeholder_url, stream=True)
-        with open(placeholder_file, "wb") as f:
-            shutil.copyfileobj(r.raw, f)
-
-    # Copy placeholder to simulate output
-    shutil.copy(placeholder_file, file_path)
+    # ---- Save frames as MP4 ----
+    clip = ImageSequenceClip(frames, fps=12)
+    clip.write_videofile(file_path, codec="libx264", audio=False, verbose=False, logger=None)
 
-    # ✅ Build full URL for Hugging Face Space
+    # Build full URL for Hugging Face Space
     base_url = "https://abrahamdw882-wan2-api.hf.space"
     video_url = f"{base_url}/file/{file_id}.mp4"
 
     return {
         "status": "success",
         "model_file": MODEL_FILE,
-        "prompt": prompt,
-        "steps": steps,
-        "video_url": video_url,
-        "note": "Replace this with actual inference code."
+        "prompt": request.prompt,
+        "steps": request.steps,
+        "video_url": video_url
     }
 
-
-# ✅ POST endpoint
-@app.post("/generate")
-def generate_post(request: PromptRequest):
-    return generate_output(request.prompt, request.steps)
-
-
-# ✅ NEW: GET endpoint with query ?q=
-@app.get("/generate")
-def generate_get(q: str = Query(..., description="Prompt for generation"), steps: int = 20):
-    return generate_output(q, steps)
-
-
-# ✅ Serve output videos
+# -------------------- Serve output videos --------------------
 app.mount("/file", StaticFiles(directory=OUTPUT_DIR), name="file")
 
-
-# ✅ Run server in HF Spaces
+# -------------------- Run server --------------------
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)