"""
Forge video-generation worker for Modal (https://modal.com).

Deploy this to YOUR OWN free Modal account (no card, $30/month recurring
compute credit, confirmed live 2026-08-23) -- Forge never sees or pays for
your Modal usage. Once deployed, Forge calls the resulting URL directly and
synchronously, the same way it already calls Hugging Face Spaces -- no
notebook session to keep running, no polling worker, no job queue.

UNVERIFIED, READ THIS FIRST:
This was written from Modal's and diffusers' own current documentation
(decorator syntax, pipeline class, checkpoint id -- all checked live
2026-08-23), not from an actual test deployment -- there is no Modal
account in the environment this was written in to deploy and run against.
Real GPU memory/timing/cost behavior can only be confirmed by actually
deploying this. If `modal deploy` or a real request fails, paste the exact
error back for a fix -- don't assume this is production-ready as-is.

Setup (one time):
  pip install modal
  modal setup                      # opens a browser to link your free account
  modal deploy modal-worker/app.py # prints your callable URL, e.g.
                                    # https://<workspace>--forge-video-generate.modal.run
Paste that URL (plus any Modal token) into Forge's "Connect your Modal
account" panel.
"""

import base64
import io

import modal

# Wan2.2-I2V-A14B is the real image-to-video checkpoint (confirmed via its
# own model card, 2026-08-23) -- a 14B-active-parameter Mixture-of-Experts
# model, meaningfully bigger than Wan2.2's smaller T2V-1.3B/TI2V-5B dense
# variants. That likely needs more VRAM than a free T4 (16GB) comfortably
# gives -- A100 is the safer default here, at the cost of burning through
# the $30/month free credit faster (roughly 12 hours/month of A100 vs. ~50
# hours/month of T4, per Modal's published per-second pricing). If real
# testing shows Wan2.2-I2V-A14B OOMs or runs fine on a T4 with lower
# precision / CPU offload, adjust GPU_TYPE below accordingly.
GPU_TYPE = "A100"
MODEL_ID = "Wan-AI/Wan2.2-I2V-A14B-Diffusers"
FPS = 16  # matches the fps convention already used by Forge's Hugging Face Spaces path
MAX_DURATION_SECONDS = 5  # conservative default -- raise once real generation time/memory is confirmed

image = (
    modal.Image.debian_slim(python_version="3.11")
    .pip_install(
        "torch",
        "git+https://github.com/huggingface/diffusers",  # Wan2.2 support needs diffusers from source, per its own model card
        "transformers",
        "accelerate",
        "imageio",
        "imageio-ffmpeg",
        "pillow",
        "fastapi[standard]",
    )
)

app = modal.App("forge-video", image=image)


@app.cls(gpu=GPU_TYPE, timeout=600, scaledown_window=300)
class Wan22Worker:
    @modal.enter()
    def load_model(self):
        # Loaded once per warm container, not on every request -- the
        # weights are several GB, reloading them per call would dominate
        # both latency and the free-credit burn rate.
        import torch
        from diffusers import WanImageToVideoPipeline

        self.pipe = WanImageToVideoPipeline.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16)
        self.pipe.to("cuda")

    @modal.fastapi_endpoint(method="POST")
    def generate(self, body: dict):
        import torch
        from diffusers.utils import export_to_video
        from PIL import Image

        image_b64 = body.get("image", "")
        prompt = body.get("prompt", "")
        duration_seconds = min(max(float(body.get("durationSeconds") or 3), 0.5), MAX_DURATION_SECONDS)
        if not image_b64:
            return {"error": "'image' is required -- this worker is image-to-video only."}
        if not prompt:
            return {"error": "'prompt' is required."}

        try:
            # Forge sends a full "data:image/...;base64,..." URI, same shape
            # its Hugging Face path already handles -- strip the prefix if present.
            raw = image_b64.split(",", 1)[1] if image_b64.startswith("data:") else image_b64
            input_image = Image.open(io.BytesIO(base64.b64decode(raw))).convert("RGB")

            num_frames = int(duration_seconds * FPS)
            with torch.inference_mode():
                frames = self.pipe(
                    image=input_image,
                    prompt=prompt,
                    num_frames=num_frames,
                    guidance_scale=5.0,
                    num_inference_steps=30,
                ).frames[0]

            out_path = "/tmp/output.mp4"
            export_to_video(frames, out_path, fps=FPS)
            with open(out_path, "rb") as f:
                video_bytes = f.read()

            return {"resultBase64": f"data:video/mp4;base64,{base64.b64encode(video_bytes).decode()}"}
        except Exception as exc:  # noqa: BLE001 -- surfaced directly to Forge's UI, real detail matters more than a clean type here
            return {"error": str(exc)}