Update logoGenertor.py

logoGenertor.py

@@ -1,72 +1,76 @@
-import torch
-from diffusers import StableDiffusionPipeline
 import gradio as gr
-import numpy as np
+from diffusers import StableDiffusionPipeline
+import torch
 from PIL import Image
+import io
 
-# Load model
+# Load the model globally once
 pipe = StableDiffusionPipeline.from_pretrained(
-    "stabilityai/stable-diffusion-2",
-    torch_dtype=torch.float16,
-    revision="fp16"
-).to("cuda" if torch.cuda.is_available() else "cpu")
+    "stabilityai/sd-turbo"
+)
+pipe.to("cpu")  # Free hardware: CPU
 
-# Style mapping
+# Style presets
 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"
+    "Minimal": "minimal, clean lines, flat design",
+    "Modern": "modern, bold, gradient",
+    "Retro": "retro, vintage, 90s style"
 }
 
-# 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)
+# Resolutions
+RESOLUTIONS = ["256x256", "384x384", "512x512", "768x768"]
 
-# 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]
+    # 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 = remove_white_background(image)
+        image = image.convert("RGBA")
 
     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")
+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)
 
-        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)
+    btn = gr.Button("Generate")
 
-        output = gr.Image(label="Generated Logo", show_download_button=True, type="pil")
-        submit = gr.Button("Generate")
+    output_img = gr.Image(label="Generated Logo", type="pil")
+    download_btn = gr.File(label="Download Logo (PNG)")
 
-        submit.click(fn=generate_logo, inputs=[prompt, style, resolution, transparent], outputs=output)
+    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")
 
-    return demo
+    btn.click(
+        fn=process_all,
+        inputs=[prompt, style, resolution, transparent, dark_mode, share_link],
+        outputs=[output_img, download_btn]
+    )
 
 if __name__ == "__main__":
-    ui = create_ui()
-    ui.launch(share=True)
+    demo.launch(share=True)