Update app.py

app.py

@@ -1,27 +1,32 @@
 import gradio as gr
 import numpy as np
 import random
-from diffusers import AutoPipelineForText2Image
+from diffusers import DiffusionPipeline
 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 = AutoPipelineForText2Image.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
+# Simplified model loading
+try:
+    if torch.cuda.is_available():
+        torch_dtype = torch.float16
+        pipe = DiffusionPipeline.from_pretrained(
+            model_repo_id,
+            torch_dtype=torch_dtype,
+            variant="fp16",
+            use_safetensors=True
+        ).to(device)
+    else:
+        torch_dtype = torch.float32
+        pipe = DiffusionPipeline.from_pretrained(
+            model_repo_id,
+            torch_dtype=torch_dtype
+        ).to(device)
+except Exception as e:
+    print(f"Error loading model: {e}")
+    raise
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
@@ -37,15 +42,15 @@ def infer(
     num_inference_steps,
     progress=gr.Progress(track_tqdm=True),
 ):
-    start_time = time()
-    
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
+    try:
+        start_time = time()
+        
+        if randomize_seed:
+            seed = random.randint(0, MAX_SEED)
 
-    generator = torch.Generator(device=device).manual_seed(seed)
+        generator = torch.Generator(device=device).manual_seed(seed)
 
-    # Generate image with progress updates
-    with torch.no_grad():
+        # Generate image
         image = pipe(
             prompt=prompt,
             negative_prompt=negative_prompt,
@@ -56,9 +61,12 @@ def infer(
             generator=generator,
         ).images[0]
 
-    gen_time = time() - start_time
-    
-    return image, seed, f"Generated in {gen_time:.2f}s"
+        gen_time = time() - start_time
+        
+        return image, seed, f"Generated in {gen_time:.2f}s"
+    except Exception as e:
+        print(f"Error during inference: {e}")
+        return None, seed, f"Error: {str(e)}"
 
 examples = [
     ["Astronaut in a jungle, cold color palette, muted colors, detailed, 8k", 1024, 1024],
@@ -68,50 +76,19 @@ examples = [
 
 css = """
 :root {
-    --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));
+    --primary: #6e6af0;
+    --secondary: #f5f5f7;
+    --accent: #f5f5f7;
+    --text: #1e1e1e;
+    --shadow: 0 4px 12px rgba(0, 0, 0, 0.1);
 }
 
 .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));
+    --primary: #a5a5fc;
+    --secondary: #2d3748;
+    --accent: #4a5568;
+    --text: #f7fafc;
+    --shadow: 0 4px 12px rgba(0, 0, 0, 0.3);
 }
 
 #col-container {
@@ -128,28 +105,21 @@ css = """
 .header h1 {
     font-size: 2.5rem;
     font-weight: 700;
-    color: var(--primary-500);
+    color: var(--primary);
     margin-bottom: 10px;
 }
 
 .header p {
-    color: var(--text-secondary);
+    color: var(--text);
+    opacity: 0.8;
 }
 
 .prompt-container, .result-container, .advanced-settings {
-    background-color: rgb(var(--surface-50));
+    background: var(--secondary);
     border-radius: 12px;
     padding: 20px;
-    box-shadow: 0 4px 12px rgba(0, 0, 0, 0.1);
+    box-shadow: var(--shadow);
     margin-bottom: 20px;
-    border: 1px solid rgb(var(--surface-200));
-}
-
-.dark .prompt-container,
-.dark .result-container,
-.dark .advanced-settings {
-    background-color: rgb(var(--surface-800));
-    border-color: rgb(var(--surface-700));
 }
 
 .advanced-settings .form {
@@ -158,10 +128,6 @@ css = """
     gap: 16px;
 }
 
-.advanced-settings .form > * {
-    margin-bottom: 0 !important;
-}
-
 .control-row {
     display: flex;
     gap: 10px;
@@ -169,13 +135,13 @@ css = """
 }
 
 .btn-primary {
-    background: var(--primary-500) !important;
+    background: var(--primary) !important;
     border: none !important;
     color: white !important;
 }
 
 .btn-primary:hover {
-    background: var(--primary-600) !important;
+    opacity: 0.9 !important;
 }
 
 .examples {
@@ -186,44 +152,23 @@ css = """
 }
 
 .example-prompt {
-    background: rgb(var(--surface-100));
+    background: var(--secondary);
     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: 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 {
-    width: 100%;
-    height: 120px;
-    object-fit: cover;
-    border-radius: 6px;
-    margin-bottom: 8px;
+    box-shadow: var(--shadow);
 }
 
-/* Theme toggle button */
 .theme-toggle {
     position: absolute;
     top: 20px;
     right: 20px;
-    background: var(--primary-100);
+    background: var(--secondary);
     border: none;
     border-radius: 50%;
     width: 40px;
@@ -232,47 +177,12 @@ css = """
     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;
 }
 """
 
@@ -294,13 +204,6 @@ document.addEventListener('DOMContentLoaded', () => {
     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') ? '☀️' : '🌙';
-        }
-    });
 });
 """
 
@@ -338,8 +241,7 @@ with gr.Blocks(css=css, js=js, theme=gr.themes.Soft()) as demo:
             result = gr.Image(
                 label="Generated Image",
                 show_label=False,
-                height=500,
-                elem_id="output-image"
+                height=500
             )
             with gr.Row():
                 seed_info = gr.Textbox(
@@ -372,7 +274,6 @@ with gr.Blocks(css=css, js=js, theme=gr.themes.Soft()) as demo:
                         randomize_seed = gr.Checkbox(
                             label="Randomize seed",
                             value=True,
-                            interactive=True
                         )
                     
                     with gr.Column():
@@ -415,8 +316,7 @@ with gr.Blocks(css=css, js=js, theme=gr.themes.Soft()) as demo:
                         examples=[[example[0], example[1], example[2]]],
                         inputs=[prompt, width, height],
                         label="",
-                        examples_per_page=20,
-                        elem_id=f"example-{example[0][:10]}"
+                        examples_per_page=20
                     )
 
     run_button.click(