Update logoGenertor.py

logoGenertor.py

@@ -1,34 +1,72 @@
-import torch
-from diffusers import StableDiffusionPipeline
-from PIL import Image
-import gradio as gr
-
-# Load the pre-trained Stable Diffusion model
-pipe = StableDiffusionPipeline.from_pretrained(
-    "runwayml/stable-diffusion-v1-5",
-    torch_dtype=torch.float16
-)
-pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu")
-
-# Logo generation function
-def generate_logo(description):
-    prompt = f"A high-quality, professional logo: {description}, sharp, modern, HD"
-    image = pipe(
-        prompt,
-        height=512,
-        width=512,
-        num_inference_steps=50,
-        guidance_scale=7.5
-    ).images[0]
-    return image.resize((1024, 1024), resample=Image.LANCZOS)
-
-# Gradio interface
-iface = gr.Interface(
-    fn=generate_logo,
-    inputs=gr.Textbox(label="Logo Description"),
-    outputs=gr.Image(type="pil", label="Generated Logo"),
-    title="AI Logo Generator",
-    description="Enter a description to generate a high-quality logo"
-)
-
-iface.launch(share=True)
+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)