Spaces:
Running
Running
| 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) | |