Update app.py

app.py

@@ -1,30 +1,85 @@
 import gradio as gr
-from diffusers import StableDiffusionPipeline
+from diffusers import StableDiffusionXLImg2ImgPipeline
+import torch
 
-# Load the Hugging Face model
-model_id = "kopyl/ui-icons-256"
-pipeline = StableDiffusionPipeline.from_pretrained(model_id)
-#pipeline.to("cuda")  # Use GPU for faster inference, if available
+# Load a lightweight pipeline that works well on CPU
+def load_image_generator():
+    try:
+        model = StableDiffusionXLImg2ImgPipeline.from_pretrained(
+            "stabilityai/stable-diffusion-xl-base-1.0", 
+            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
 
-def generate_icon(prompt):
-    """Generate an icon from a text prompt."""
-    result = pipeline(prompt).images[0]  # Generate the image
-    return result
-
-# Build Gradio Interface
-with gr.Blocks() as app:
-    gr.Markdown("## 🖌️ Icon Generator App")
-    gr.Markdown("Generate 256x256 icons using the `kopyl/ui-icons-256` model.")
-    
-    with gr.Row():
-        prompt_input = gr.Textbox(label="Enter a prompt", placeholder="E.g., 'a colorful app icon for a weather app'")
-        generate_button = gr.Button("Generate Icon")
+# 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
     
-    output_image = gr.Image(label="Generated Icon")
+    # Default icon if no initial image
+    default_init_image = torch.randn((1, 3, 512, 512))
     
-    def on_click(prompt):
-        return generate_icon(prompt)
+    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]
+        
+        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")
+        
+        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")
+        
+        # Connect generate button to function
+        generate_btn.click(
+            fn=generate_chatbot_icon, 
+            inputs=[prompt], 
+            outputs=[output_image]
+        )
     
-    generate_button.click(on_click, inputs=prompt_input, outputs=output_image)
+    return demo
 
-app.launch()  # Launch the app
+# Launch the app
+if __name__ == "__main__":
+    demo = create_gradio_interface()
+    demo.launch()