Spaces:
Running
Running
File size: 7,115 Bytes
06ae90e 057604e cfcc0d9 06ae90e cfcc0d9 06ae90e cfcc0d9 06ae90e cfcc0d9 057604e cfcc0d9 057604e cfcc0d9 057604e cfcc0d9 057604e cfcc0d9 06ae90e cfcc0d9 06ae90e 057604e cfcc0d9 057604e cfcc0d9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 |
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() |