import torch
from diffusers import StableDiffusionImg2ImgPipeline
from PIL import Image
from io import BytesIO

class ImageToImage:
    """
    Class to handle Image-to-Image transformations using Stable Diffusion.
    """
    def __init__(self, device="cpu"):
        # Model and repository details
        model_id = "OFA-Sys/small-stable-diffusion-v0"
        self.pipe = StableDiffusionImg2ImgPipeline.from_pretrained(model_id, torch_dtype=torch.float32).pipe.to(device)
        self.device = device
        print("Image-to-Image model loaded successfully.")
    
    
    async def transform_image(self, image, prompt):
        """
        Transform an uploaded image based on a text prompt.

        Args:
            image (PIL.Image): The input image to transform.
            prompt (str): The text prompt to guide the transformation.

        Returns:
            PIL.Image: The transformed image.
        """
        if not prompt:
            raise ValueError("Prompt cannot be empty.")

        # Resize the image as required by the model
        init_image = image.resize((512, 512))
        with torch.no_grad():
            transformed_image = self.pipe(
                prompt=prompt, 
                image=init_image, 
                strength=0.75, 
                guidance_scale=7.5
            ).images[0]
        return transformed_image