GENAI_Project / logoGenertor.py
farazify's picture
Update logoGenertor.py
fb14b89 verified
import gradio as gr
from diffusers import StableDiffusionPipeline
import torch
from PIL import Image
import io
# Load the model globally once
pipe = StableDiffusionPipeline.from_pretrained(
"stabilityai/sd-turbo"
)
pipe.to("cpu") # Free hardware: CPU
# Style presets
STYLE_PRESETS = {
"Minimal": "minimal, clean lines, flat design",
"Modern": "modern, bold, gradient",
"Retro": "retro, vintage, 90s style"
}
# Resolutions
RESOLUTIONS = ["256x256", "384x384", "512x512", "768x768"]
def generate_logo(prompt, style, resolution, transparent):
# Apply style
style_prompt = STYLE_PRESETS.get(style, "")
full_prompt = f"{style_prompt}, {prompt}".strip()
# Resolution
width, height = map(int, resolution.split("x"))
# Run pipeline
image = pipe(full_prompt, height=height, width=width, num_inference_steps=15).images[0]
# Convert to RGBA if transparent
if transparent:
image = image.convert("RGBA")
return image
def get_image_bytes(img):
buffer = io.BytesIO()
img.save(buffer, format="PNG")
buffer.seek(0)
return buffer
with gr.Blocks(theme=gr.themes.Default(), title="🎨 AI Logo Generator") as demo:
gr.Markdown("## 🎨 AI Logo Generator\nEnter a description to generate a high-quality logo")
with gr.Row():
prompt = gr.Textbox(label="Logo Description", placeholder="e.g., A minimalist logo of owl for tech startup")
style = gr.Dropdown(label="Style", choices=list(STYLE_PRESETS.keys()), value="Minimal")
resolution = gr.Dropdown(label="Resolution", choices=RESOLUTIONS, value="512x512")
with gr.Row():
transparent = gr.Checkbox(label="Make background transparent", value=False)
dark_mode = gr.Checkbox(label="Dark Mode UI", value=False)
share_link = gr.Checkbox(label="Make app public (Gradio share link)", value=True)
btn = gr.Button("Generate")
output_img = gr.Image(label="Generated Logo", type="pil")
download_btn = gr.File(label="Download Logo (PNG)")
def process_all(prompt, style, resolution, transparent, dark_mode, share_link):
image = generate_logo(prompt, style, resolution, transparent)
file = get_image_bytes(image)
return image, (file, "logo.png")
btn.click(
fn=process_all,
inputs=[prompt, style, resolution, transparent, dark_mode, share_link],
outputs=[output_img, download_btn]
)
if __name__ == "__main__":
demo.launch(share=True)