Update app.py

app.py

@@ -5,10 +5,13 @@ from PIL import Image, ImageDraw, ImageFont
 import torch
 from diffusers import DiffusionPipeline
 import io
+import time
 
 # ===== CONFIG =====
 device = "cuda" if torch.cuda.is_available() else "cpu"
 torch_dtype = torch.float16 if device == "cuda" else torch.float32
+
+# Using SDXL Turbo for fastest generation
 model_repo_id = "stabilityai/sdxl-turbo"
 pipe = DiffusionPipeline.from_pretrained(
     model_repo_id,
@@ -17,108 +20,161 @@ pipe = DiffusionPipeline.from_pretrained(
 )
 pipe.to(device)
 
+# Enable memory efficient attention and channels last for better performance
+pipe.enable_xformers_memory_efficient_attention()
+pipe.unet.to(memory_format=torch.channels_last)
+
 MAX_SEED = np.iinfo(np.int32).max
-IMAGE_WIDTH = 768
-IMAGE_HEIGHT = 768
+IMAGE_SIZE = 1024  # Same as original code
 WATERMARK_TEXT = "SelamGPT"
 
-# ===== WATERMARK FUNCTION =====
+# ===== OPTIMIZED WATERMARK FUNCTION =====
 def add_watermark(image):
-    draw = ImageDraw.Draw(image)
-    font_size = int(image.width * 0.03)
+    """Optimized watermark function matching original style"""
     try:
-        font = ImageFont.truetype("Roboto-Bold.ttf", font_size)
-    except:
-        font = ImageFont.load_default()
-    text_width = draw.textlength(WATERMARK_TEXT, font=font)
-    x = image.width - text_width - 12
-    y = image.height - font_size - 10
-    draw.text((x+1, y+1), WATERMARK_TEXT, font=font, fill=(0, 0, 0, 128))
-    draw.text((x, y), WATERMARK_TEXT, font=font, fill=(255, 255, 255))
-    return image
-
-# ===== INFERENCE FUNCTION =====
+        draw = ImageDraw.Draw(image)
+        font_size = 24  # Fixed size as in original
+        
+        try:
+            font = ImageFont.truetype("Roboto-Bold.ttf", font_size)
+        except:
+            font = ImageFont.load_default(font_size)
+        
+        text_width = draw.textlength(WATERMARK_TEXT, font=font)
+        x = image.width - text_width - 10
+        y = image.height - 34
+        
+        # Shadow effect
+        draw.text((x+1, y+1), WATERMARK_TEXT, font=font, fill=(0, 0, 0, 128))
+        draw.text((x, y), WATERMARK_TEXT, font=font, fill=(255, 255, 255))
+        
+        return image
+    except Exception as e:
+        print(f"Watermark error: {str(e)}")
+        return image
+
+# ===== ULTRA-FAST INFERENCE FUNCTION =====
 def generate(
     prompt,
-    negative_prompt,
-    seed,
-    randomize_seed,
-    guidance_scale,
-    num_inference_steps,
+    negative_prompt="",
+    seed=None,
+    randomize_seed=True,
+    guidance_scale=0.0,  # 0.0 for turbo models
+    num_inference_steps=1,  # Can be as low as 1-2 for turbo
     progress=gr.Progress(track_tqdm=True),
 ):
     if not prompt.strip():
         return None, "⚠️ Please enter a prompt"
 
-    if randomize_seed:
+    start_time = time.time()
+    
+    # Seed handling
+    if randomize_seed or seed is None:
         seed = random.randint(0, MAX_SEED)
-
+    
     generator = torch.manual_seed(seed)
+    
+    # Ultra-fast generation with minimal steps
     result = pipe(
         prompt=prompt,
         negative_prompt=negative_prompt,
-        width=IMAGE_WIDTH,
-        height=IMAGE_HEIGHT,
+        width=IMAGE_SIZE,
+        height=IMAGE_SIZE,
         guidance_scale=guidance_scale,
-        num_inference_steps=num_inference_steps,
+        num_inference_steps=max(1, num_inference_steps),  # Minimum 1 step
         generator=generator,
     ).images[0]
-
+    
+    # Optimized watermark and JPG conversion
     watermarked = add_watermark(result)
     buffer = io.BytesIO()
-    watermarked.convert("RGB").save(buffer, format="JPEG", quality=70)
+    watermarked.save(buffer, format="JPEG", quality=85, optimize=True)
     buffer.seek(0)
-    return Image.open(buffer), seed
+    
+    gen_time = time.time() - start_time
+    status = f"✔️ Generated in {gen_time:.2f}s | Seed: {seed}"
+    
+    return Image.open(buffer), status
 
 # ===== EXAMPLES =====
 examples = [
-    "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
-    "An astronaut riding a green horse",
-    "A delicious ceviche cheesecake slice",
+    ["An ancient Aksumite warrior in cyberpunk armor, 4k detailed"],
+    ["Traditional Ethiopian coffee ceremony in zero gravity"],
+    ["Portrait of a Habesha queen with golden jewelry"]
 ]
 
-# ===== INTERFACE =====
-css = "#container { max-width: 700px; margin: auto; }"
-
-with gr.Blocks(css=css, title="SelamGPT Turbo Generator") as demo:
-    with gr.Column(elem_id="container"):
-        gr.Markdown("# 🖼️ SelamGPT Image Generator")
+# ===== OPTIMIZED INTERFACE =====
+theme = gr.themes.Default(
+    primary_hue="emerald",
+    secondary_hue="amber",
+    font=[gr.themes.GoogleFont("Poppins"), "Arial", "sans-serif"]
+)
 
-        with gr.Row():
+with gr.Blocks(theme=theme, title="SelamGPT Turbo Generator") as demo:
+    gr.Markdown("""
+    # 🎨 SelamGPT Turbo Image Generator
+    *Ultra-fast 1024x1024 image generation with SDXL-Turbo*
+    """)
+    
+    with gr.Row():
+        with gr.Column(scale=3):
             prompt = gr.Textbox(
-                label="Prompt",
-                show_label=False,
-                placeholder="Enter your prompt",
-                lines=1,
-                scale=3
+                label="Describe your image",
+                placeholder="A futuristic Ethiopian city with flying cars...",
+                lines=3,
+                max_lines=5
             )
-            generate_btn = gr.Button("Generate", variant="primary")
-
-        output_image = gr.Image(label="Generated Image", type="pil", format="jpeg")
-        seed_display = gr.Textbox(label="Seed Used", interactive=False)
-
-        with gr.Accordion("⚙️ Advanced Settings", open=False):
-            negative_prompt = gr.Textbox(label="Negative Prompt", placeholder="What to avoid (optional)", max_lines=1)
-            randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
-            seed = gr.Slider(0, MAX_SEED, step=1, label="Seed", value=0)
-
-            guidance_scale = gr.Slider(0.0, 10.0, step=0.1, label="Guidance Scale", value=0.0)
-            num_inference_steps = gr.Slider(1, 10, step=1, label="Inference Steps", value=2)
-
-        gr.Examples(examples=examples, inputs=[prompt])
-
-        generate_btn.click(
-            fn=generate,
-            inputs=[
-                prompt,
-                negative_prompt,
-                seed,
-                randomize_seed,
-                guidance_scale,
-                num_inference_steps
-            ],
-            outputs=[output_image, seed_display]
+            with gr.Row():
+                generate_btn = gr.Button("Generate Image", variant="primary")
+                clear_btn = gr.Button("Clear")
+            
+            gr.Examples(
+                examples=examples,
+                inputs=[prompt]
+            )
+            
+        with gr.Column(scale=2):
+            output_image = gr.Image(
+                label="Generated Image",
+                type="pil",
+                format="jpeg",
+                height=512
+            )
+            status_output = gr.Textbox(
+                label="Status",
+                interactive=False
+            )
+    
+    with gr.Accordion("⚙️ Advanced Settings", open=False):
+        negative_prompt = gr.Textbox(
+            label="Negative Prompt",
+            placeholder="What to avoid (optional)",
+            max_lines=1
         )
+        with gr.Row():
+            randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
+            seed = gr.Number(label="Seed", value=0, precision=0)
+        guidance_scale = gr.Slider(0.0, 1.0, value=0.0, step=0.1, label="Guidance Scale")
+        num_inference_steps = gr.Slider(1, 4, value=1, step=1, label="Inference Steps")
+
+    generate_btn.click(
+        fn=generate,
+        inputs=[
+            prompt,
+            negative_prompt,
+            seed,
+            randomize_seed,
+            guidance_scale,
+            num_inference_steps
+        ],
+        outputs=[output_image, status_output]
+    )
+    
+    clear_btn.click(
+        fn=lambda: [None, ""],
+        outputs=[output_image, status_output]
+    )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.queue(max_size=4)  # Increased queue for better throughput
+    demo.launch(server_name="0.0.0.0", server_port=7860)