Spaces:
Running
Running
| import gradio as gr | |
| import torch | |
| from diffusers import StableDiffusionXLImg2ImgPipeline | |
| from PIL import Image, ImageOps | |
| 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() | |
| 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 | |
| def round_image_corners(self, image, corner_radius): | |
| # Create a rounded corner mask | |
| if corner_radius == 0: | |
| return image | |
| # Create a new image with an alpha channel | |
| rounded_image = Image.new('RGBA', image.size, (0, 0, 0, 0)) | |
| # Create a mask for rounded corners | |
| mask = Image.new('L', image.size, 255) | |
| from PIL import ImageDraw | |
| draw = ImageDraw.Draw(mask) | |
| # 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 | |
| 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__": | |
| main() |