File size: 4,574 Bytes
15115fa b2a588a 15115fa b2a588a 15115fa b2a588a 15115fa b2a588a 15115fa b2a588a 15115fa b2a588a 15115fa b2a588a 15115fa b2a588a 15115fa b2a588a 15115fa b2a588a 15115fa b2a588a 15115fa b2a588a 15115fa b2a588a 15115fa b2a588a 15115fa b2a588a 15115fa b2a588a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 | """
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") |