import torch
from diffusers import StableDiffusionPipeline
from PIL import Image
from io import BytesIO
from utils import load_unet_model

class TextToImage:
    """
    Class to handle Text-to-Image generation using Stable Diffusion XL.
    """
    def __init__(self, device="cpu"):
        # Model and repository details
        model_id = "OFA-Sys/small-stable-diffusion-v0"
        self.device = device
        self.pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32).pipe.to(device)
        print("Text-to-Image model loaded successfully.")
    
    
    async def generate_image(self, prompt):
        """
        Generate an image from a text prompt.

        Args:
            prompt (str): The text prompt to generate the image.

        Returns:
            PIL.Image: The generated image.
        """
        with torch.no_grad():
            image = self.pipe(
                prompt
            ).images[0]
        return image