|
from typing import cast, Union |
|
|
|
import torch |
|
|
|
from diffusers import AutoencoderKLWan |
|
from diffusers.video_processor import VideoProcessor |
|
from diffusers.utils import export_to_video |
|
|
|
|
|
class EndpointHandler: |
|
def __init__(self, path=""): |
|
self.device = "cuda" |
|
self.dtype = torch.float16 |
|
self.vae = cast( |
|
AutoencoderKLWan, |
|
AutoencoderKLWan.from_pretrained(path, torch_dtype=self.dtype) |
|
.to(self.device, self.dtype) |
|
.eval(), |
|
) |
|
|
|
self.vae_scale_factor_temporal = ( |
|
2 ** sum(self.vae.temperal_downsample) if getattr(self, "vae", None) else 4 |
|
) |
|
self.vae_scale_factor_spatial = ( |
|
2 ** len(self.vae.temperal_downsample) if getattr(self, "vae", None) else 8 |
|
) |
|
self.video_processor = VideoProcessor( |
|
vae_scale_factor=self.vae_scale_factor_spatial |
|
) |
|
|
|
@torch.no_grad() |
|
def __call__(self, data) -> Union[torch.Tensor, bytes]: |
|
""" |
|
Args: |
|
data (:obj:): |
|
includes the input data and the parameters for the inference. |
|
""" |
|
tensor = cast(torch.Tensor, data["inputs"]) |
|
parameters = cast(dict, data.get("parameters", {})) |
|
do_scaling = cast(bool, parameters.get("do_scaling", True)) |
|
output_type = cast(str, parameters.get("output_type", "pil")) |
|
partial_postprocess = cast(bool, parameters.get("partial_postprocess", False)) |
|
if partial_postprocess and output_type != "pt": |
|
output_type = "pt" |
|
|
|
tensor = tensor.to(self.device, self.dtype) |
|
|
|
if do_scaling: |
|
latents_mean = ( |
|
torch.tensor(self.vae.config.latents_mean) |
|
.view(1, self.vae.config.z_dim, 1, 1, 1) |
|
.to(tensor.device, tensor.dtype) |
|
) |
|
latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view( |
|
1, self.vae.config.z_dim, 1, 1, 1 |
|
).to(tensor.device, tensor.dtype) |
|
tensor = tensor / latents_std + latents_mean |
|
|
|
with torch.no_grad(): |
|
frames = cast(torch.Tensor, self.vae.decode(tensor, return_dict=False)[0]) |
|
|
|
if partial_postprocess: |
|
frames = frames[0].permute(1, 0, 2, 3) |
|
frames = torch.stack([(frame * 0.5 + 0.5).clamp(0, 1) for frame in frames]) |
|
frames = frames.permute(0, 2, 3, 1).contiguous().float() |
|
frames = (frames * 255).round().to(torch.uint8) |
|
elif output_type == "pil": |
|
frames = cast( |
|
torch.Tensor, |
|
self.video_processor.postprocess_video(frames, output_type="pt")[0], |
|
) |
|
elif output_type == "mp4": |
|
frames = cast( |
|
torch.Tensor, |
|
self.video_processor.postprocess_video(frames, output_type="pil")[0], |
|
) |
|
path = export_to_video(frames, fps=16) |
|
with open(path, "rb") as f: |
|
frames = f.read() |
|
elif output_type == "pt": |
|
frames = frames |
|
|
|
return frames |
|
|
