Update app.py

app.py

@@ -1,30 +1,35 @@
 import torch
-import spaces
 import gradio as gr
 from diffusers import DiffusionPipeline
 
-# Load the pipeline once at startup
+# ================== DEVICE ==================
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"Using device: {device}")
+
+# ================== LOAD PIPELINE ==================
 print("Loading Z-Image-Turbo pipeline...")
 pipe = DiffusionPipeline.from_pretrained(
     "Tongyi-MAI/Z-Image-Turbo",
-    torch_dtype=torch.bfloat16,
     low_cpu_mem_usage=False,
 )
-pipe.to("cuda")
-
-# ======== AoTI compilation + FA3 ========
-# pipe.transformer.layers._repeated_blocks = ["ZImageTransformerBlock"]
-# spaces.aoti_blocks_load(pipe.transformer.layers, "zerogpu-aoti/Z-Image", variant="fa3")
-
+pipe = pipe.to(device)
 print("Pipeline loaded!")
 
-@spaces.GPU
-def generate_image(prompt, height, width, num_inference_steps, seed, randomize_seed, progress=gr.Progress(track_tqdm=True)):
-    """Generate an image from the given prompt."""
+# ================== IMAGE GENERATION ==================
+def generate_image(
+    prompt,
+    height,
+    width,
+    num_inference_steps,
+    seed,
+    randomize_seed,
+    progress=gr.Progress(track_tqdm=True),
+):
     if randomize_seed:
         seed = torch.randint(0, 2**32 - 1, (1,)).item()
-    
-    generator = torch.Generator("cuda").manual_seed(int(seed))
+
+    generator = torch.Generator(device).manual_seed(int(seed))
+
     image = pipe(
         prompt=prompt,
         height=int(height),
@@ -33,19 +38,20 @@ def generate_image(prompt, height, width, num_inference_steps, seed, randomize_s
         guidance_scale=0.0,
         generator=generator,
     ).images[0]
-    
+
     return image, seed
 
-# Example prompts
+
+# ================== EXAMPLES ==================
 examples = [
-    ["Young Chinese woman in red Hanfu, intricate embroidery. Impeccable makeup, red floral forehead pattern. Elaborate high bun, golden phoenix headdress, red flowers, beads. Holds round folding fan with lady, trees, bird. Neon lightning-bolt lamp, bright yellow glow, above extended left palm. Soft-lit outdoor night background, silhouetted tiered pagoda, blurred colorful distant lights."],
-    ["A majestic dragon soaring through clouds at sunset, scales shimmering with iridescent colors, detailed fantasy art style"],
-    ["Cozy coffee shop interior, warm lighting, rain on windows, plants on shelves, vintage aesthetic, photorealistic"],
-    ["Astronaut riding a horse on Mars, cinematic lighting, sci-fi concept art, highly detailed"],
-    ["Portrait of a wise old wizard with a long white beard, holding a glowing crystal staff, magical forest background"],
+    ["Young Chinese woman in red Hanfu, intricate embroidery, neon lightning lamp floating above palm, cinematic lighting"],
+    ["A majestic dragon soaring through clouds at sunset, fantasy art, ultra detailed"],
+    ["Cozy coffee shop interior, rain on windows, warm light, photorealistic"],
+    ["Astronaut riding a horse on Mars, cinematic sci-fi concept art"],
+    ["Portrait of an old wizard with glowing staff, magical forest"],
 ]
 
-# Custom theme with modern aesthetics (Gradio 6)
+# ================== THEME ==================
 custom_theme = gr.themes.Soft(
     primary_hue="yellow",
     secondary_hue="amber",
@@ -53,206 +59,85 @@ custom_theme = gr.themes.Soft(
     font=gr.themes.GoogleFont("Inter"),
     text_size="lg",
     spacing_size="md",
-    radius_size="lg"
-).set(
-    button_primary_background_fill="*primary_500",
-    button_primary_background_fill_hover="*primary_600",
-    block_title_text_weight="600",
+    radius_size="lg",
 )
 
-# Build the Gradio interface
-with gr.Blocks(fill_height=True) as demo:
-    # Header
+# ================== UI ==================
+with gr.Blocks(fill_height=True, theme=custom_theme) as demo:
     gr.Markdown(
         """
-        # burak image
-        **Ultra-fast AI image generation** • Generate stunning images in just 8 steps
-        """,
-        elem_classes="header-text"
+        # 🤖 Burak Image
+        **Ultra-fast AI image generation** • CPU / GPU Auto
+        """
     )
-    
-    with gr.Row(equal_height=False):
-        # Left column - Input controls
+
+    with gr.Row():
         with gr.Column(scale=1, min_width=320):
             prompt = gr.Textbox(
-                label="✨ Your Prompt",
-                placeholder="Describe the image you want to create...",
+                label="✨ Prompt",
+                placeholder="Describe the image you want...",
                 lines=5,
-                max_lines=10,
-                autofocus=True,
             )
-            
+
             with gr.Accordion("⚙️ Advanced Settings", open=False):
                 with gr.Row():
-                    height = gr.Slider(
-                        minimum=512,
-                        maximum=2048,
-                        value=1024,
-                        step=64,
-                        label="Height",
-                        info="Image height in pixels"
-                    )
-                    width = gr.Slider(
-                        minimum=512,
-                        maximum=2048,
-                        value=1024,
-                        step=64,
-                        label="Width",
-                        info="Image width in pixels"
-                    )
-                
+                    height = gr.Slider(512, 2048, 1024, step=64, label="Height")
+                    width = gr.Slider(512, 2048, 1024, step=64, label="Width")
+
                 num_inference_steps = gr.Slider(
-                    minimum=1,
-                    maximum=20,
-                    value=9,
-                    step=1,
-                    label="Inference Steps",
-                    info="9 steps = 8 DiT forwards (recommended)"
+                    1, 20, 9, step=1, label="Inference Steps"
                 )
-                
-                with gr.Row():
-                    randomize_seed = gr.Checkbox(
-                        label="🎲 Random Seed",
-                        value=True,
-                    )
-                    seed = gr.Number(
-                        label="Seed",
-                        value=42,
-                        precision=0,
-                        visible=False,
-                    )
-                
-                def toggle_seed(randomize):
-                    return gr.Number(visible=not randomize)
-                
+
+                randomize_seed = gr.Checkbox(
+                    label="🎲 Random Seed", value=True
+                )
+                seed = gr.Number(
+                    label="Seed", value=42, precision=0, visible=False
+                )
+
                 randomize_seed.change(
-                    toggle_seed,
-                    inputs=[randomize_seed],
-                    outputs=[seed]
+                    lambda x: gr.Number(visible=not x),
+                    randomize_seed,
+                    seed,
                 )
-            
+
             generate_btn = gr.Button(
                 "🚀 Generate Image",
                 variant="primary",
                 size="lg",
-                scale=1
             )
-            
-            # Example prompts
+
             gr.Examples(
                 examples=examples,
                 inputs=[prompt],
-                label="💡 Try these prompts",
-                examples_per_page=5,
+                label="💡 Example Prompts",
             )
-        
-        # Right column - Output
+
         with gr.Column(scale=1, min_width=320):
             output_image = gr.Image(
                 label="Generated Image",
                 type="pil",
-                show_label=False,
                 height=600,
                 buttons=["download", "share"],
             )
-            
+
             used_seed = gr.Number(
                 label="🎲 Seed Used",
                 interactive=False,
-                container=True,
             )
-    
-    # Footer credits
-    gr.Markdown(
-        """
-        ---
-        <div style="text-align: center; opacity: 0.7; font-size: 0.9em; margin-top: 1rem;">
-        <strong>Model:</strong> <a href="https://huggingface.co/Tongyi-MAI/Z-Image-Turbo" target="_blank">Tongyi-MAI/Z-Image-Turbo</a> (Apache 2.0 License) • 
-        <strong>Demo by:</strong> <a href="https://x.com/realmrfakename" target="_blank">@mrfakename</a> • 
-        <strong>Redesign by:</strong> AnyCoder • 
-        <strong>Optimizations:</strong> <a href="https://huggingface.co/multimodalart" target="_blank">@multimodalart</a> (FA3 + AoTI)
-        </div>
-        """,
-        elem_classes="footer-text"
-    )
-    
-    # Connect the generate button
+
     generate_btn.click(
-        fn=generate_image,
+        generate_image,
         inputs=[prompt, height, width, num_inference_steps, seed, randomize_seed],
         outputs=[output_image, used_seed],
     )
-    
-    # Also allow generating by pressing Enter in the prompt box
+
     prompt.submit(
-        fn=generate_image,
+        generate_image,
         inputs=[prompt, height, width, num_inference_steps, seed, randomize_seed],
         outputs=[output_image, used_seed],
     )
 
+# ================== LAUNCH ==================
 if __name__ == "__main__":
-    demo.launch(
-        theme=custom_theme,
-        css="""
-        .header-text h1 {
-            font-size: 2.5rem !important;
-            font-weight: 700 !important;
-            margin-bottom: 0.5rem !important;
-            background: linear-gradient(135deg, #fbbf24 0%, #f59e0b 100%);
-            -webkit-background-clip: text;
-            -webkit-text-fill-color: transparent;
-            background-clip: text;
-        }
-        
-        .header-text p {
-            font-size: 1.1rem !important;
-            color: #64748b !important;
-            margin-top: 0 !important;
-        }
-        
-        .footer-text {
-            padding: 1rem 0;
-        }
-        
-        .footer-text a {
-            color: #f59e0b !important;
-            text-decoration: none !important;
-            font-weight: 500;
-        }
-        
-        .footer-text a:hover {
-            text-decoration: underline !important;
-        }
-        
-        /* Mobile optimizations */
-        @media (max-width: 768px) {
-            .header-text h1 {
-                font-size: 1.8rem !important;
-            }
-            
-            .header-text p {
-                font-size: 1rem !important;
-            }
-        }
-        
-        /* Smooth transitions */
-        button, .gr-button {
-            transition: all 0.2s ease !important;
-        }
-        
-        button:hover, .gr-button:hover {
-            transform: translateY(-1px);
-            box-shadow: 0 4px 12px rgba(0, 0, 0, 0.15) !important;
-        }
-        
-        /* Better spacing */
-        .gradio-container {
-            max-width: 1400px !important;
-            margin: 0 auto !important;
-        }
-        """,
-        footer_links=[
-            "api",
-            "gradio"
-        ]
-    )
+    demo.launch()