Update app.py

app.py

@@ -1,7 +1,7 @@
 import gradio as gr
+from diffusers import DiffusionPipeline
 import torch
-from diffusers import StableDiffusionXLImg2ImgPipeline
-from PIL import Image, ImageOps
+from PIL import Image, ImageDraw
 
 class ChatbotIconGenerator:
     def __init__(self):
@@ -15,25 +15,16 @@ class ChatbotIconGenerator:
         
         # Predefined prompt templates
         self.PROMPT_TEMPLATES = [
-            # Professional/Corporate
-            "Professional AI chatbot avatar, minimalist design, sleek geometric shapes, corporate blue and white color palette",
-            "Elegant corporate chatbot icon, modern flat design, clean lines, subtle technology motif",
-            
-            # Cute/Friendly
             "Cute cartoon chatbot mascot, big eyes, friendly smile, pastel colors, kawaii style",
+            "Professional AI chatbot avatar, minimalist design, sleek geometric shapes, corporate blue and white",
+            "Futuristic AI chatbot icon, glowing circuit patterns, metallic blue and silver",
+            "Abstract geometric chatbot avatar, clean lines, single color gradient background",
+            "Watercolor style chatbot icon, soft brush strokes, dreamy color blend",
             "Adorable robot character avatar, round shape, soft colors, playful expression",
-            
-            # Sci-Fi/Tech
-            "Futuristic AI chatbot icon, glowing circuit patterns, metallic blue and silver, high-tech aesthetic",
-            "Cyberpunk chatbot avatar, neon accents, digital glitch effects, modern tech design",
-            
-            # Minimalist
             "Ultra-minimalist chatbot icon, simple geometric face, monochrome color scheme",
-            "Abstract geometric chatbot avatar, clean lines, single color gradient background",
-            
-            # Artistic
-        "Watercolor style chatbot icon, soft brush strokes, dreamy color blend, artistic interpretation",
-            "Sketch-style chatbot avatar, hand-drawn look, pencil texture, artistic rendering"
+            "Cyberpunk chatbot avatar, neon accents, digital glitch effects",
+            "Elegant corporate chatbot icon, modern flat design, clean lines",
+            "Sketch-style chatbot avatar, hand-drawn look, pencil texture"
         ]
         
         # Rounding options
@@ -49,11 +40,12 @@ class ChatbotIconGenerator:
     
     def load_image_generator(self):
         try:
-            model = StableDiffusionXLImg2ImgPipeline.from_pretrained(
-                "stabilityai/stable-diffusion-2-1", 
+            # Use a more lightweight model
+            model = DiffusionPipeline.from_pretrained(
+                "runwayml/stable-diffusion-v1-5", # kopyl/ui-icons-256
                 torch_dtype=torch.float16,
-                variant="fp16",
-                use_safetensors=True
+                safety_checker=None,  # Disable safety checker to reduce load
+                requires_safety_checker=False
             )
             return model.to("cpu")
         except Exception as e:
@@ -70,7 +62,6 @@ class ChatbotIconGenerator:
         
         # Create a mask for rounded corners
         mask = Image.new('L', image.size, 255)
-        from PIL import ImageDraw
         draw = ImageDraw.Draw(mask)
         
         # Draw rounded rectangle
@@ -89,26 +80,22 @@ class ChatbotIconGenerator:
         prompt, 
         size, 
         corner_rounding,
-        negative_prompt="low quality, bad composition, blurry, ugly",
-        num_inference_steps=25,
-        guidance_scale=8.0,
-        strength=0.75
+        negative_prompt="low quality, bad composition, blurry, ugly, deformed",
+        num_inference_steps=20,
+        guidance_scale=7.5
     ):
         if self.model is None:
-            return None
+            raise ValueError("Model failed to load. Please check your dependencies.")
         
         try:
-            # Create a random initial image of specified size
-            default_init_image = torch.randn((1, 3, size, size))
-            
             # Generate the image
             generated_image = self.model(
                 prompt=prompt,
                 negative_prompt=negative_prompt,
                 num_inference_steps=num_inference_steps,
                 guidance_scale=guidance_scale,
-                strength=strength,
-                image=default_init_image
+                height=size,
+                width=size
             ).images[0]
             
             # Resize and round corners
@@ -120,7 +107,7 @@ class ChatbotIconGenerator:
         
         except Exception as e:
             print(f"Error generating image: {e}")
-            return None
+            raise
     
     def create_gradio_interface(self):
         with gr.Blocks(title="🤖 Chatbot Icon Generator") as demo: