File size: 1,910 Bytes
77a6e14 | 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 | import base64
import io
import torch
from diffusers import AutoencoderKLWan, WanPipeline
from diffusers.utils import export_to_video
class EndpointHandler:
def __init__(self, path=""):
self.device = "cuda" if torch.cuda.is_available() else "cpu"
self.dtype = torch.bfloat16
# Load VAE and pipeline from the diffusers-format repo
vae = AutoencoderKLWan.from_pretrained(
"Wan-AI/Wan2.2-TI2V-5B-Diffusers",
subfolder="vae",
torch_dtype=torch.float32,
)
self.pipe = WanPipeline.from_pretrained(
"Wan-AI/Wan2.2-TI2V-5B-Diffusers",
vae=vae,
torch_dtype=self.dtype,
)
self.pipe.to(self.device)
def __call__(self, data):
prompt = data.pop("inputs", "")
params = data.pop("parameters", {})
num_frames = params.get("num_frames", 49)
height = params.get("height", 480)
width = params.get("width", 832)
num_inference_steps = params.get("num_inference_steps", 30)
fps = params.get("fps", 24)
guidance_scale = params.get("guidance_scale", 5.0)
negative_prompt = params.get("negative_prompt", "low quality, blurry, distorted")
frames = self.pipe(
prompt=prompt,
negative_prompt=negative_prompt,
num_frames=num_frames,
height=height,
width=width,
num_inference_steps=num_inference_steps,
guidance_scale=guidance_scale,
).frames[0]
# Export to video in memory
buf = io.BytesIO()
export_to_video(frames, buf, fps=fps)
video_b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
duration = len(frames) / fps
return {
"video": video_b64,
"format": "mp4",
"frames": len(frames),
"duration": duration,
}
|