generator.py

from diffusers import AutoPipelineForText2Image
from PIL import Image
import numpy as np
import cv2

def load_model():
    model_id = "stabilityai/sd-turbo"
    pipe = AutoPipelineForText2Image.from_pretrained(model_id)
    return pipe

def generate_image(pipe, prompt, steps=8, guidance=2.5, width=512, height=512):
    image = pipe(
        prompt=prompt,
        num_inference_steps=steps,
        guidance_scale=guidance,
        width=width,
        height=height
    ).images[0]
    return image

def upscale_image(pil_image, scale=2):
    img = np.array(pil_image)
    upscaled = cv2.resize(img, None, fx=scale, fy=scale, interpolation=cv2.INTER_CUBIC)
    return Image.fromarray(upscaled)