Update app.py

app.py

@@ -1,85 +1,195 @@
 import gradio as gr
-from diffusers import StableDiffusionXLImg2ImgPipeline
 import torch
+from diffusers import StableDiffusionXLImg2ImgPipeline
+from PIL import Image, ImageOps
 
-# Load a lightweight pipeline that works well on CPU
-def load_image_generator():
-    try:
-        model = StableDiffusionXLImg2ImgPipeline.from_pretrained(
-            "stabilityai/stable-diffusion-2-1", 
-            torch_dtype=torch.float16,
-            variant="fp16",
-            use_safetensors=True
-        )
-        # Ensure it runs on CPU
-        #model = model.to("cpu")
-        return model
-    except Exception as e:
-        print(f"Error loading model: {e}")
-        return None
-
-# Generate chatbot icon
-def generate_chatbot_icon(
-    prompt, 
-    negative_prompt="low quality, bad composition, blurry", 
-    num_inference_steps=20,
-    guidance_scale=7.5,
-    strength=0.75
-):
-    # Load the model
-    model = load_image_generator()
-    if model is None:
-        return None
+class ChatbotIconGenerator:
+    def __init__(self):
+        # Predefined size options
+        self.SIZE_OPTIONS = {
+            "Small (128x128)": 128,
+            "Medium (256x256)": 256,
+            "Large (512x512)": 512,
+            "Extra Large (1024x1024)": 1024
+        }
+        
+        # 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",
+            "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"
+        ]
+        
+        # Rounding options
+        self.CORNER_OPTIONS = {
+            "No Rounding": 0,
+            "Slight Rounding": 20,
+            "Medium Rounding": 50,
+            "Full Rounded": 100
+        }
+        
+        # Load the model
+        self.model = self.load_image_generator()
     
-    # Default icon if no initial image
-    default_init_image = torch.randn((1, 3, 512, 512))
+    def load_image_generator(self):
+        try:
+            model = StableDiffusionXLImg2ImgPipeline.from_pretrained(
+                "stabilityai/stable-diffusion-2-1", 
+                torch_dtype=torch.float16,
+                variant="fp16",
+                use_safetensors=True
+            )
+            return model.to("cpu")
+        except Exception as e:
+            print(f"Error loading model: {e}")
+            return None
     
-    try:
-        # Generate the image
-        image = model(
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            num_inference_steps=num_inference_steps,
-            guidance_scale=guidance_scale,
-            strength=strength,
-            image=default_init_image
-        ).images[0]
+    def round_image_corners(self, image, corner_radius):
+        # Create a rounded corner mask
+        if corner_radius == 0:
+            return image
         
-        return image
-    except Exception as e:
-        print(f"Error generating image: {e}")
-        return None
-
-# Create Gradio interface
-def create_gradio_interface():
-    with gr.Blocks() as demo:
-        gr.Markdown("# 🤖 Chatbot Icon Generator")
+        # Create a new image with an alpha channel
+        rounded_image = Image.new('RGBA', image.size, (0, 0, 0, 0))
         
-        with gr.Row():
-            with gr.Column():
-                # Prompt input
-                prompt = gr.Textbox(
-                    label="Icon Description", 
-                    value="Cute minimalist chatbot avatar, clean design, friendly expression, cartoon style"
-                )
-                
-                # Generate button
-                generate_btn = gr.Button("Generate Icon")
-                
-            with gr.Column():
-                # Output image
-                output_image = gr.Image(label="Generated Chatbot Icon")
+        # Create a mask for rounded corners
+        mask = Image.new('L', image.size, 255)
+        from PIL import ImageDraw
+        draw = ImageDraw.Draw(mask)
         
-        # Connect generate button to function
-        generate_btn.click(
-            fn=generate_chatbot_icon, 
-            inputs=[prompt], 
-            outputs=[output_image]
+        # Draw rounded rectangle
+        draw.rounded_rectangle(
+            [0, 0, image.width-1, image.height-1], 
+            radius=corner_radius, 
+            fill=255
         )
+        
+        # Paste the original image with the mask
+        rounded_image.paste(image, mask=mask)
+        return rounded_image
+    
+    def generate_chatbot_icon(
+        self, 
+        prompt, 
+        size, 
+        corner_rounding,
+        negative_prompt="low quality, bad composition, blurry, ugly",
+        num_inference_steps=25,
+        guidance_scale=8.0,
+        strength=0.75
+    ):
+        if self.model is None:
+            return None
+        
+        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
+            ).images[0]
+            
+            # Resize and round corners
+            generated_image = generated_image.resize((size, size))
+            rounded_image = self.round_image_corners(generated_image, 
+                self.CORNER_OPTIONS[corner_rounding])
+            
+            return rounded_image
+        
+        except Exception as e:
+            print(f"Error generating image: {e}")
+            return None
     
-    return demo
+    def create_gradio_interface(self):
+        with gr.Blocks(title="🤖 Chatbot Icon Generator") as demo:
+            gr.Markdown("# 🤖 Chatbot Icon Generator")
+            
+            with gr.Row():
+                with gr.Column():
+                    # Prompt selection
+                    prompt_dropdown = gr.Dropdown(
+                        label="Quick Templates", 
+                        choices=self.PROMPT_TEMPLATES,
+                        allow_custom_value=True
+                    )
+                    
+                    # Custom prompt input
+                    custom_prompt = gr.Textbox(
+                        label="Custom Prompt (Optional)", 
+                        placeholder="Enter your own detailed description..."
+                    )
+                    
+                    # Size selection
+                    size_dropdown = gr.Dropdown(
+                        label="Icon Size", 
+                        choices=list(self.SIZE_OPTIONS.keys()),
+                        value="Medium (256x256)"
+                    )
+                    
+                    # Corner rounding
+                    corner_dropdown = gr.Dropdown(
+                        label="Corner Rounding", 
+                        choices=list(self.CORNER_OPTIONS.keys()),
+                        value="Slight Rounding"
+                    )
+                    
+                    # Generate button
+                    generate_btn = gr.Button("Generate Icon", variant="primary")
+                
+                with gr.Column():
+                    # Output image
+                    output_image = gr.Image(label="Generated Chatbot Icon")
+            
+            # Logic for prompt selection
+            def update_prompt(template):
+                return template
+            
+            prompt_dropdown.change(
+                fn=update_prompt, 
+                inputs=[prompt_dropdown], 
+                outputs=[custom_prompt]
+            )
+            
+            # Generate button logic
+            generate_btn.click(
+                fn=lambda prompt, size, corners: self.generate_chatbot_icon(
+                    prompt or "Cute minimalist chatbot avatar, clean design, friendly expression",
+                    self.SIZE_OPTIONS[size],
+                    corners
+                ),
+                inputs=[custom_prompt, size_dropdown, corner_dropdown], 
+                outputs=[output_image]
+            )
+        
+        return demo
 
 # Launch the app
+def main():
+    generator = ChatbotIconGenerator()
+    demo = generator.create_gradio_interface()
+    demo.launch()
+
 if __name__ == "__main__":
-    demo = create_gradio_interface()
-    demo.launch()
+    main()