handler.py

"""
handler.py
----------
Handler personalizado para un Inference Endpoint dedicado de Hugging Face.

Este fichero se sube AL REPO DEL MODELO (no al Space), en la raiz, junto
a los pesos. Cuando creas el Endpoint desde la web de HF, el sistema
detecta este handler.py y lo usa en lugar de intentar mapear el modelo
a una tarea estandar.

Requiere que el repo tenga tambien un requirements.txt con, como minimo:
    diffusers>=0.31.0
    transformers
    accelerate
    imageio[ffmpeg]
    torch

Nota: la clase real del pipeline de Wan (WanImageToVideoPipeline o el
nombre que tenga en tu version de diffusers) puede variar segun cuando
la instales. Comprueba el nombre exacto en la documentacion de diffusers
antes de desplegar, porque el soporte de Wan es reciente y puede cambiar.
"""

import base64
import io
import tempfile

import ftfy
import torch
from PIL import Image
from diffusers import WanImageToVideoPipeline
from diffusers.utils import export_to_video

# Bug conocido de transformers: el tokenizer de CLIP importa ftfy de forma
# local dentro de __init__, pero luego lo referencia como nombre global en
# otros metodos (p.ej. bpe()), lo que provoca NameError: name 'ftfy' is not
# defined aunque el paquete este instalado. Se corrige inyectando el modulo
# como atributo global antes de que se cargue el pipeline.
import transformers.models.clip.tokenization_clip as _clip_tokenization_module
_clip_tokenization_module.ftfy = ftfy


class EndpointHandler:
    """Handler que Hugging Face Inference Endpoints instancia una vez al arrancar."""

    def __init__(self, path: str = ""):
        """
        Carga el pipeline en GPU al arrancar el contenedor.

        path es la ruta local donde el Endpoint descarga los pesos del repo.
        """
        self.pipe = WanImageToVideoPipeline.from_pretrained(
            path,
            torch_dtype=torch.bfloat16,
        )
        self.pipe.to("cuda")

        # El VAE y el image_encoder suelen necesitar float32 por estabilidad
        # numerica en estos pipelines de imagen-a-video; forzarlos aqui evita
        # el mismatch de dtype con el resto del pipeline en bfloat16.
        if getattr(self.pipe, "vae", None) is not None:
            self.pipe.vae = self.pipe.vae.to(torch.float32)
        if getattr(self.pipe, "image_encoder", None) is not None:
            self.pipe.image_encoder = self.pipe.image_encoder.to(torch.float32)

    def __call__(self, data: dict) -> dict:
        """
        Procesa una peticion de generacion de video.

        data llega con la forma {"inputs": {...}} tal y como lo envia
        el cliente REST. Devuelve un dict serializable en JSON con el
        video codificado en base64.
        """
        inputs = data.get("inputs", data)

        imagen_b64 = inputs["image"]
        prompt = inputs.get("prompt", "")
        negative_prompt = inputs.get("negative_prompt", "")
        num_frames = inputs.get("num_frames", 33)
        num_inference_steps = inputs.get("num_inference_steps", 20)
        guidance_scale = inputs.get("guidance_scale", 5.0)
        seed = inputs.get("seed", 42)
        height = inputs.get("height", 480)
        width = inputs.get("width", 832)
        fps = inputs.get("fps", 16)

        imagen = decodificar_imagen(imagen_b64)
        generador = torch.Generator(device="cuda").manual_seed(seed)

        salida = self.pipe(
            image=imagen,
            prompt=prompt,
            negative_prompt=negative_prompt,
            num_frames=num_frames,
            num_inference_steps=num_inference_steps,
            guidance_scale=guidance_scale,
            height=height,
            width=width,
            generator=generador,
        )

        frames = salida.frames[0]
        video_b64 = exportar_video_base64(frames, fps)

        return {"video_base64": video_b64}


def decodificar_imagen(imagen_b64: str) -> Image.Image:
    """Convierte una imagen codificada en base64 a un objeto PIL en modo RGB."""
    imagen_bytes = base64.b64decode(imagen_b64)
    return Image.open(io.BytesIO(imagen_bytes)).convert("RGB")


def exportar_video_base64(frames, fps: int) -> str:
    """Exporta una lista de frames a mp4 temporal y devuelve el contenido en base64."""
    with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmp:
        export_to_video(frames, tmp.name, fps=fps)
        tmp.seek(0)
        video_bytes = tmp.read()

    return base64.b64encode(video_bytes).decode("utf-8")