Delete text_to_image.py

text_to_image.py

@@ -1,56 +0,0 @@
-import torch
-import spaces
-from diffusers import StableDiffusionXLPipeline, EulerDiscreteScheduler
-from PIL import Image
-from io import BytesIO
-from utils import load_unet_model
-
-@spaces.GPU
-class TextToImage:
-    """
-    Class to handle Text-to-Image generation using Stable Diffusion XL.
-    """
-    def __init__(self, device="cuda"):
-        # 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 Text-to-Image model...")
-        self.unet = load_unet_model(self.base, self.repo, self.ckpt, device=self.device)
-        
-        # Initialize the pipeline
-        self.pipe = StableDiffusionXLPipeline.from_pretrained(
-            self.base, 
-            unet=self.unet, 
-            torch_dtype=torch.float16, 
-            variant="fp16"
-        ).to(self.device)
-        
-        # Set the scheduler
-        self.pipe.scheduler = EulerDiscreteScheduler.from_config(
-            self.pipe.scheduler.config, 
-            timestep_spacing="trailing"
-        )
-        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, 
-                num_inference_steps=4, 
-                guidance_scale=0
-            ).images[0]
-        return image