Spaces:
Running
Running
File size: 2,607 Bytes
59c2540 |
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 |
import torch
from diffusers import StableDiffusionPipeline
import gradio as gr
import numpy as np
from PIL import Image
# Load model
pipe = StableDiffusionPipeline.from_pretrained(
"stabilityai/stable-diffusion-2",
torch_dtype=torch.float16,
revision="fp16"
).to("cuda" if torch.cuda.is_available() else "cpu")
# Style mapping
STYLE_PRESETS = {
"Minimal": "minimal flat logo, vector style",
"Modern": "modern logo, professional design",
"Retro": "vintage retro logo, 90s style",
"Tech": "technology startup logo, digital design",
"Luxury": "luxury brand logo, elegant and premium"
}
# Convert to transparent PNG if background is white-ish
def remove_white_background(pil_image):
image = pil_image.convert("RGBA")
data = np.array(image)
r, g, b, a = data.T
white_areas = (r > 240) & (g > 240) & (b > 240)
data[..., :-1][white_areas.T] = (255, 255, 255)
data[..., -1][white_areas.T] = 0 # transparent
return Image.fromarray(data)
# Generate image
def generate_logo(prompt, style, resolution, transparent):
if not prompt:
return None
styled_prompt = f"{prompt}, {STYLE_PRESETS[style]}"
width, height = map(int, resolution.split('x'))
image = pipe(styled_prompt, height=height, width=width).images[0]
if transparent:
image = remove_white_background(image)
return image
# Build UI
def create_ui():
with gr.Blocks(theme=gr.themes.Soft(), css=".dark-mode body { background-color: #1e1e1e; color: white; }") as demo:
gr.Markdown("# 🎨 AI Logo Generator")
gr.Markdown("Enter a description to generate a high-quality logo")
with gr.Row():
prompt = gr.Textbox(label="Logo Description", placeholder="e.g., A minimalist owl logo for a tech startup")
style = gr.Dropdown(choices=list(STYLE_PRESETS.keys()), label="Style", value="Minimal")
resolution = gr.Dropdown(choices=["512x512", "768x768"], label="Resolution", value="512x512")
with gr.Row():
transparent = gr.Checkbox(label="Make background transparent", value=True)
dark_mode = gr.Checkbox(label="Dark Mode UI", value=False)
share_space = gr.Checkbox(label="Make app public (Gradio share link)", value=False)
output = gr.Image(label="Generated Logo", show_download_button=True, type="pil")
submit = gr.Button("Generate")
submit.click(fn=generate_logo, inputs=[prompt, style, resolution, transparent], outputs=output)
return demo
if __name__ == "__main__":
ui = create_ui()
ui.launch(share=True)
|