logo-generator / handler.py
jonathanrstern's picture
more logs
f53f2d2
raw
history blame
689 Bytes
from typing import Dict, List, Any
from diffusers import ControlNetModel, StableDiffusionControlNetPipeline, StableDiffusionPipeline, DiffusionPipeline
class EndpointHandler():
def __init__(self, path=""):
self.pipeline = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0")
self.pipeline.load_lora_weights(".", weight_name="/repository/safetensors.safetensors")
print('hiiiii')
def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
print('data:', data)
input = data['inputs']
print('input:', input)
result = self.pipeline(input)
print('result:', result)
return result