handler.py

import torch
from diffusers import StableDiffusionXLPipeline
import base64
from io import BytesIO

class InferenceHandler:
    def __init__(self):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        model_name = "colt12/maxcushion"

        # If your model is private, include the use_auth_token parameter
        self.pipe = StableDiffusionXLPipeline.from_pretrained(
            model_name,
            torch_dtype=torch.float16,
            use_safetensors=True,
            # Uncomment the line below and replace with your token if needed
            # use_auth_token="your_huggingface_token"
        ).to(self.device)

    def __call__(self, inputs):
        prompt = inputs.get("prompt", "")
        if not prompt:
            raise ValueError("A prompt must be provided")
        
        negative_prompt = inputs.get("negative_prompt", "")
        
        image = self.pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            num_inference_steps=30,
            guidance_scale=7.5
        ).images[0]
        
        buffered = BytesIO()
        image.save(buffered, format="PNG")
        image_base64 = base64.b64encode(buffered.getvalue()).decode("utf-8")
        
        return {"image_base64": image_base64}

handler = InferenceHandler()