| from typing import Dict, List, Any | |
| from diffusers import DiffusionPipeline | |
| import torch | |
| import diffusers | |
| class EndpointHandler(): | |
| def __init__(self, path=""): | |
| self.path = path | |
| self.model = "remg1997/dynabench-sdxl10" | |
| self.pipeline = DiffusionPipeline.from_pretrained(self.model, torch_dtype=torch.float16, use_safetensors=True, variant="fp16") | |
| self.pipeline = self.pipeline.to("cuda", torch.float16) | |
| def __call__(self, data: Dict[str, Any])-> List[Dict[str, Any]]: | |
| print("Torch version is", torch.__version__) | |
| print("Diffusers version is", diffusers.__version__) | |
| inputs = data.pop("inputs", data) | |
| print("inputs", inputs) | |
| steps = data.pop("steps", 30) | |
| image = self.pipeline(inputs, num_inference_steps = steps) | |
| return [{"image": image}] | |