| import torch | |
| from diffusers import DiffusionPipeline | |
| class EndpointHandler: | |
| def __init__(self, model_dir): | |
| self.pipe = DiffusionPipeline.from_pretrained( | |
| model_dir, | |
| custom_pipeline="pipeline_wan_i2v", | |
| torch_dtype=torch.float16 | |
| ).to("cuda") | |
| def __call__(self, data): | |
| image = data["image"] | |
| output = self.pipe(image) | |
| return output |