| 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}] |
|
|
|
|
|
|
