from diffusers import StableDiffusionPipeline
import torch

def init():
    global model
    model_id = "colt12/maxcushion"
    model = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
    model.to("cuda")

def inference(model_inputs):
    global model
    prompt = model_inputs.get('prompt', '')
    image = model(prompt).images[0]
    # Convert image to bytes
    import io
    buf = io.BytesIO()
    image.save(buf, format='PNG')
    byte_im = buf.getvalue()
    return {"image_bytes": byte_im}