Update image_to_image.py

image_to_image.py

@@ -1,60 +1,41 @@
-import torch
-from diffusers import StableDiffusionXLImg2ImgPipeline, EulerDiscreteScheduler
-from PIL import Image
-from io import BytesIO
-from utils import load_unet_model
-
-class ImageToImage:
-    """
-    Class to handle Image-to-Image transformations using Stable Diffusion XL.
-    """
-    def __init__(self, device="cpu"):
-        # Model and repository details
-        self.base = "stabilityai/stable-diffusion-xl-base-1.0"
-        self.repo = "ByteDance/SDXL-Lightning"
-        self.ckpt = "sdxl_lightning_4step_unet.safetensors"
-        self.device = device
-
-        # Load the UNet model
-        print("Loading Image-to-Image model...")
-        self.unet = load_unet_model(self.base, self.repo, self.ckpt, device=self.device)
-        
-        # Initialize the pipeline
-        self.pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained(
-            self.base, 
-            unet=self.unet, 
-            torch_dtype=torch.float32
-        ).to(self.device)
-        
-        # Set the scheduler
-        self.pipe.scheduler = EulerDiscreteScheduler.from_config(
-            self.pipe.scheduler.config, 
-            timestep_spacing="trailing"
-        )
-        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((768, 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
+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