app.py

# app.py
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

app = FastAPI(title="WAN2 GGUF API", version="1.0")

# ✅ Directories
MODEL_REPO = "calcuis/wan2-gguf"
MODEL_FILE = "wan2.2-animate-14b-q4_0.gguf"  # big model
MODEL_DIR = "models"
OUTPUT_DIR = "outputs"

os.makedirs(MODEL_DIR, exist_ok=True)
os.makedirs(OUTPUT_DIR, exist_ok=True)

# ✅ Download model file at startup
model_path = hf_hub_download(
    repo_id=MODEL_REPO,
    filename=MODEL_FILE,
    local_dir=MODEL_DIR
)
print("✅ Model downloaded to:", model_path)


# Request schema
class PromptRequest(BaseModel):
    prompt: str
    steps: int = 20


@app.get("/")
def root():
    return {"message": "WAN2 GGUF API is running!"}


@app.post("/generate")
def generate_video(request: PromptRequest):
    """
    Dummy video generator — for now just copies a placeholder .mp4.
    Replace this later with actual WAN2 inference code.
    """

    # Unique filename
    file_id = str(uuid.uuid4())
    file_path = os.path.join(OUTPUT_DIR, f"{file_id}.mp4")

    # Use a Hugging Face placeholder video
    placeholder_url = (
        "https://huggingface.co/datasets/huggingface/documentation-images/"
        "resolve/main/video-placeholder.mp4"
    )

    # Download placeholder only once
    placeholder_file = os.path.join(OUTPUT_DIR, "placeholder.mp4")
    if not os.path.exists(placeholder_file):
        import requests
        r = requests.get(placeholder_url, stream=True)
        with open(placeholder_file, "wb") as f:
            shutil.copyfileobj(r.raw, f)

    # Copy placeholder to simulate unique output
    shutil.copy(placeholder_file, file_path)

    return {
        "status": "success",
        "model_file": MODEL_FILE,
        "prompt": request.prompt,
        "steps": request.steps,
        "video_url": f"/file/{file_id}.mp4",
        "note": "Replace this with actual inference code."
    }


# ✅ Serve output videos
app.mount("/file", StaticFiles(directory=OUTPUT_DIR), name="file")


# ✅ Run server in HF Spaces
if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)