Update app.py

app.py

@@ -1,19 +1,27 @@
 import gradio as gr
 import numpy as np
 import random
-from diffusers import DiffusionPipeline
+from diffusers import AutoPipelineForText2Image
 import torch
+from time import time
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model_repo_id = "stabilityai/sdxl-turbo"
 
+# Load model with optimizations
 if torch.cuda.is_available():
     torch_dtype = torch.float16
+    pipe = AutoPipelineForText2Image.from_pretrained(
+        model_repo_id, 
+        torch_dtype=torch_dtype,
+        variant="fp16",
+        use_safetensors=True
+    )
+    pipe.enable_model_cpu_offload()  # For better memory management
+    pipe.enable_xformers_memory_efficient_attention()  # Faster attention
 else:
     torch_dtype = torch.float32
-
-pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
-pipe = pipe.to(device)
+    pipe = AutoPipelineForText2Image.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
@@ -29,22 +37,28 @@ def infer(
     num_inference_steps,
     progress=gr.Progress(track_tqdm=True),
 ):
+    start_time = time()
+    
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
 
-    generator = torch.Generator().manual_seed(seed)
+    generator = torch.Generator(device=device).manual_seed(seed)
 
-    image = pipe(
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        guidance_scale=guidance_scale,
-        num_inference_steps=num_inference_steps,
-        width=width,
-        height=height,
-        generator=generator,
-    ).images[0]
+    # Generate image with progress updates
+    with torch.no_grad():
+        image = pipe(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            guidance_scale=guidance_scale,
+            num_inference_steps=num_inference_steps,
+            width=width,
+            height=height,
+            generator=generator,
+        ).images[0]
 
-    return image, seed
+    gen_time = time() - start_time
+    
+    return image, seed, f"Generated in {gen_time:.2f}s"
 
 examples = [
     ["Astronaut in a jungle, cold color palette, muted colors, detailed, 8k", 1024, 1024],
@@ -54,11 +68,50 @@ examples = [
 
 css = """
 :root {
-    --primary: #6e6af0;
-    --secondary: #f5f5f7;
-    --accent: #f5f5f7;
-    --text: #1e1e1e;
-    --shadow: 0 4px 12px rgba(0, 0, 0, 0.1);
+    --primary-50: #f0f0ff;
+    --primary-100: #e0e0ff;
+    --primary-200: #c7c7fe;
+    --primary-300: #a5a5fc;
+    --primary-400: #8181f8;
+    --primary-500: #6e6af0;
+    --primary-600: #5a56e4;
+    --primary-700: #4a46c9;
+    --primary-800: #3e3ba3;
+    --primary-900: #383682;
+    --primary-950: #211f4d;
+    
+    --surface-0: 255 255 255;
+    --surface-50: 248 250 252;
+    --surface-100: 241 245 249;
+    --surface-200: 226 232 240;
+    --surface-300: 203 213 225;
+    --surface-400: 148 163 184;
+    --surface-500: 100 116 139;
+    --surface-600: 71 85 105;
+    --surface-700: 45 55 72;
+    --surface-800: 30 41 59;
+    --surface-900: 15 23 42;
+    --surface-950: 3 6 23;
+    
+    --text-primary: rgb(var(--surface-900));
+    --text-secondary: rgb(var(--surface-600));
+}
+
+.dark {
+    --primary-50: #211f4d;
+    --primary-100: #383682;
+    --primary-200: #3e3ba3;
+    --primary-300: #4a46c9;
+    --primary-400: #5a56e4;
+    --primary-500: #6e6af0;
+    --primary-600: #8181f8;
+    --primary-700: #a5a5fc;
+    --primary-800: #c7c7fe;
+    --primary-900: #e0e0ff;
+    --primary-950: #f0f0ff;
+    
+    --text-primary: rgb(var(--surface-100));
+    --text-secondary: rgb(var(--surface-300));
 }
 
 #col-container {
@@ -75,36 +128,28 @@ css = """
 .header h1 {
     font-size: 2.5rem;
     font-weight: 700;
-    color: var(--primary);
+    color: var(--primary-500);
     margin-bottom: 10px;
 }
 
 .header p {
-    color: var(--text);
-    opacity: 0.8;
+    color: var(--text-secondary);
 }
 
-.prompt-container {
-    background: white;
+.prompt-container, .result-container, .advanced-settings {
+    background-color: rgb(var(--surface-50));
     border-radius: 12px;
     padding: 20px;
-    box-shadow: var(--shadow);
+    box-shadow: 0 4px 12px rgba(0, 0, 0, 0.1);
     margin-bottom: 20px;
+    border: 1px solid rgb(var(--surface-200));
 }
 
-.result-container {
-    background: white;
-    border-radius: 12px;
-    padding: 20px;
-    box-shadow: var(--shadow);
-    margin-bottom: 20px;
-}
-
-.advanced-settings {
-    background: white;
-    border-radius: 12px;
-    padding: 20px;
-    box-shadow: var(--shadow);
+.dark .prompt-container,
+.dark .result-container,
+.dark .advanced-settings {
+    background-color: rgb(var(--surface-800));
+    border-color: rgb(var(--surface-700));
 }
 
 .advanced-settings .form {
@@ -124,12 +169,13 @@ css = """
 }
 
 .btn-primary {
-    background: var(--primary) !important;
+    background: var(--primary-500) !important;
     border: none !important;
+    color: white !important;
 }
 
 .btn-primary:hover {
-    opacity: 0.9 !important;
+    background: var(--primary-600) !important;
 }
 
 .examples {
@@ -140,16 +186,28 @@ css = """
 }
 
 .example-prompt {
-    background: var(--secondary);
+    background: rgb(var(--surface-100));
     padding: 12px;
     border-radius: 8px;
     cursor: pointer;
     transition: all 0.2s;
+    border: 1px solid rgb(var(--surface-200));
+}
+
+.dark .example-prompt {
+    background: rgb(var(--surface-700));
+    border-color: rgb(var(--surface-600));
 }
 
 .example-prompt:hover {
-    background: #e0e0e8;
+    background: var(--primary-100);
     transform: translateY(-2px);
+    border-color: var(--primary-300);
+}
+
+.dark .example-prompt:hover {
+    background: var(--primary-800);
+    border-color: var(--primary-600);
 }
 
 .example-img {
@@ -160,14 +218,93 @@ css = """
     margin-bottom: 8px;
 }
 
+/* Theme toggle button */
+.theme-toggle {
+    position: absolute;
+    top: 20px;
+    right: 20px;
+    background: var(--primary-100);
+    border: none;
+    border-radius: 50%;
+    width: 40px;
+    height: 40px;
+    display: flex;
+    align-items: center;
+    justify-content: center;
+    cursor: pointer;
+    transition: all 0.2s;
+}
+
+.theme-toggle:hover {
+    background: var(--primary-200);
+}
+
+.dark .theme-toggle {
+    background: var(--primary-800);
+}
+
+.dark .theme-toggle:hover {
+    background: var(--primary-700);
+}
+
 @media (max-width: 768px) {
     .advanced-settings .form {
         grid-template-columns: 1fr;
     }
+    
+    .theme-toggle {
+        top: 10px;
+        right: 10px;
+    }
+}
+
+/* Loading animation */
+@keyframes spin {
+    0% { transform: rotate(0deg); }
+    100% { transform: rotate(360deg); }
+}
+
+.loading-spinner {
+    display: inline-block;
+    width: 20px;
+    height: 20px;
+    border: 3px solid rgba(255, 255, 255, 0.3);
+    border-radius: 50%;
+    border-top-color: white;
+    animation: spin 1s ease-in-out infinite;
+    margin-right: 8px;
 }
 """
 
-with gr.Blocks(css=css) as demo:
+js = """
+function toggleTheme() {
+    const body = document.body;
+    body.classList.toggle('dark');
+    localStorage.setItem('gradio-theme', body.classList.contains('dark') ? 'dark' : 'light');
+}
+
+document.addEventListener('DOMContentLoaded', () => {
+    const savedTheme = localStorage.getItem('gradio-theme') || 'light';
+    if (savedTheme === 'dark') {
+        document.body.classList.add('dark');
+    }
+    
+    const themeToggle = document.createElement('button');
+    themeToggle.className = 'theme-toggle';
+    themeToggle.innerHTML = savedTheme === 'dark' ? '☀️' : '🌙';
+    themeToggle.onclick = toggleTheme;
+    document.body.appendChild(themeToggle);
+    
+    // Update icon when theme changes
+    document.body.addEventListener('click', (e) => {
+        if (e.target === themeToggle) {
+            themeToggle.innerHTML = document.body.classList.contains('dark') ? '☀️' : '🌙';
+        }
+    });
+});
+"""
+
+with gr.Blocks(css=css, js=js, theme=gr.themes.Soft()) as demo:
     with gr.Column(elem_id="col-container"):
         with gr.Column(visible=True) as header:
             gr.Markdown(
@@ -204,11 +341,15 @@ with gr.Blocks(css=css) as demo:
                 height=500,
                 elem_id="output-image"
             )
-            seed_info = gr.Textbox(
-                label="Seed",
-                interactive=False,
-                visible=False
-            )
+            with gr.Row():
+                seed_info = gr.Textbox(
+                    label="Seed",
+                    interactive=False
+                )
+                time_info = gr.Textbox(
+                    label="Generation Time",
+                    interactive=False
+                )
 
         with gr.Accordion("🛠️ Advanced Settings", open=False, elem_classes="advanced-settings"):
             with gr.Column(elem_classes="form"):
@@ -290,7 +431,7 @@ with gr.Blocks(css=css) as demo:
             guidance_scale,
             num_inference_steps,
         ],
-        outputs=[result, seed_info],
+        outputs=[result, seed_info, time_info],
     )
     prompt.submit(
         fn=infer,
@@ -304,8 +445,8 @@ with gr.Blocks(css=css) as demo:
             guidance_scale,
             num_inference_steps,
         ],
-        outputs=[result, seed_info],
+        outputs=[result, seed_info, time_info],
     )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.queue(api_open=False).launch()