Update app.py

app.py

@@ -2,59 +2,32 @@ import spaces
 import gradio as gr
 import torch
 from diffusers import ZImagePipeline
+from io import BytesIO
 import os
 from pathlib import Path
 
 # Load the model directly at startup
 print("Loading Z-Image Turbo model...")
-print("This may take a few minutes on first run while the model downloads...")
-
-
-# Load the pipeline with optimal settings
 pipe = ZImagePipeline.from_pretrained(
     "Tongyi-MAI/Z-Image-Turbo",
     torch_dtype=torch.bfloat16,
     low_cpu_mem_usage=False,
 )
-    
-# Move to GPU if available
 device = "cuda" if torch.cuda.is_available() else "cpu"
 pipe.to(device)
 print(f"Model loaded on {device}")
-    
-print("Model loaded successfully!")
 
 @spaces.GPU()
-def generate_image(
-    prompt,
-    progress=gr.Progress(track_tqdm=True)
-):
-    """
-    Generate an image using Z-Image Turbo model.
-    
-    Args:
-        prompt: Text description of the desired image
-    
-    Returns:
-        Generated PIL Image
-    """
-    global pipe
-    
+def generate_image(prompt, progress=gr.Progress(track_tqdm=True)):
     if pipe is None:
         raise gr.Error("Model failed to load on startup. Please restart the application.")
-    
     if not prompt.strip():
         raise gr.Error("Please enter a prompt to generate an image.")
     
-    # Determine device
     device = "cuda" if torch.cuda.is_available() else "cpu"
-    
-    # Set random seed for reproducibility
     generator = torch.Generator(device).manual_seed(42)
-    
-    # Generate the image with optimal settings
     progress(0.1, desc="Generating image...")
-    
+
     try:
         result = pipe(
             prompt=prompt,
@@ -65,296 +38,28 @@ def generate_image(
             guidance_scale=0.0,
             generator=generator,
         )
-        
         image = result.images[0]
+
+        # Convert to PNG in memory
+        buffer = BytesIO()
+        image.save(buffer, format="PNG")
+        buffer.seek(0)
         progress(1.0, desc="Complete!")
-        
-        return image
-    
+        return buffer  # Return PNG bytes
+
     except Exception as e:
         raise gr.Error(f"Generation failed: {str(e)}")
 
-# Apple-style CSS
-apple_css = """
-/* Global Styles */
-.gradio-container {
-    max-width: 980px !important;
-    margin: 0 auto !important;
-    padding: 48px 20px !important;
-    font-family: -apple-system, BlinkMacSystemFont, 'Inter', 'Segoe UI', 'Roboto', sans-serif !important;
-}
-
-/* Header */
-.header-container {
-    text-align: center;
-    margin-bottom: 48px;
-}
-
-.main-title {
-    font-size: 56px !important;
-    font-weight: 600 !important;
-    letter-spacing: -0.02em !important;
-    line-height: 1.07 !important;
-    color: #1d1d1f !important;
-    margin: 0 0 12px 0 !important;
-}
-
-.subtitle {
-    font-size: 21px !important;
-    font-weight: 400 !important;
-    line-height: 1.38 !important;
-    color: #6e6e73 !important;
-    margin: 0 0 24px 0 !important;
-}
-
-.attribution-link {
-    display: inline-block;
-    font-size: 14px !important;
-    color: #0071e3 !important;
-    text-decoration: none !important;
-    font-weight: 400 !important;
-    transition: color 0.2s ease !important;
-}
-
-.attribution-link:hover {
-    color: #0077ed !important;
-    text-decoration: underline !important;
-}
-
-/* Input Section */
-.input-section {
-    background: #ffffff;
-    border-radius: 18px;
-    padding: 32px;
-    margin-bottom: 24px;
-    box-shadow: 0 2px 12px rgba(0, 0, 0, 0.08);
-}
-
-/* Textbox */
-textarea {
-    font-size: 17px !important;
-    line-height: 1.47 !important;
-    border-radius: 12px !important;
-    border: 1px solid #d2d2d7 !important;
-    padding: 12px 16px !important;
-    transition: all 0.2s ease !important;
-    background: #ffffff !important;
-    font-family: -apple-system, BlinkMacSystemFont, 'Inter', sans-serif !important;
-}
-
-textarea:focus {
-    border-color: #0071e3 !important;
-    box-shadow: 0 0 0 4px rgba(0, 113, 227, 0.15) !important;
-    outline: none !important;
-}
-
-textarea::placeholder {
-    color: #86868b !important;
-}
-
-/* Button */
-button.primary {
-    font-size: 17px !important;
-    font-weight: 400 !important;
-    padding: 12px 32px !important;
-    border-radius: 980px !important;
-    background: #0071e3 !important;
-    border: none !important;
-    color: #ffffff !important;
-    min-height: 44px !important;
-    transition: all 0.2s ease !important;
-    letter-spacing: -0.01em !important;
-    cursor: pointer !important;
-}
-
-button.primary:hover {
-    background: #0077ed !important;
-    transform: scale(1.02) !important;
-}
+# Gradio interface (simplified)
+with gr.Blocks(title="Z-Image Turbo", fill_height=False) as demo:
+    gr.HTML("<h1>Z-Image Turbo</h1><p>Generate PNG images from prompts</p>")
+    prompt_input = gr.Textbox(placeholder="Describe the image...", lines=3)
+    generate_btn = gr.Button("Generate PNG")
+    output_image = gr.Image(type="file", label="Generated Image", show_label=False, container=False)
 
-button.primary:active {
-    transform: scale(0.98) !important;
-}
-
-/* Output Section */
-.output-section {
-    background: #ffffff;
-    border-radius: 18px;
-    padding: 32px;
-    box-shadow: 0 2px 12px rgba(0, 0, 0, 0.08);
-    overflow: hidden;
-}
-
-.output-section img {
-    border-radius: 12px !important;
-    width: 100% !important;
-    height: auto !important;
-}
-
-/* Footer */
-.footer-text {
-    text-align: center;
-    margin-top: 48px;
-    font-size: 14px !important;
-    color: #86868b !important;
-    line-height: 1.43 !important;
-}
-
-/* Progress */
-.progress-bar {
-    background: #0071e3 !important;
-    border-radius: 4px !important;
-}
-
-/* Dark Mode */
-.dark .main-title {
-    color: #f5f5f7 !important;
-}
-
-.dark .subtitle {
-    color: #a1a1a6 !important;
-}
-
-.dark .input-section,
-.dark .output-section {
-    background: #1d1d1f;
-    box-shadow: 0 2px 12px rgba(0, 0, 0, 0.4);
-}
-
-.dark textarea {
-    background: #1d1d1f !important;
-    border-color: #424245 !important;
-    color: #f5f5f7 !important;
-}
-
-.dark textarea::placeholder {
-    color: #86868b !important;
-}
-
-/* Responsive */
-@media (max-width: 734px) {
-    .main-title {
-        font-size: 40px !important;
-    }
-    
-    .subtitle {
-        font-size: 19px !important;
-    }
-    
-    .gradio-container {
-        padding: 32px 16px !important;
-    }
-    
-    .input-section,
-    .output-section {
-        padding: 24px !important;
-    }
-}
-
-/* Remove default Gradio styling */
-.contain {
-    padding: 0 !important;
-}
-"""
-
-# Create the interface
-with gr.Blocks(
-    title="Z-Image Turbo",
-    fill_height=False,
-) as demo:
-    
-    # Header
-    gr.HTML("""
-        <div class="header-container">
-            <h1 class="main-title">Z-Image Turbo</h1>
-            <p class="subtitle">Transform your ideas into stunning visuals with AI</p>
-            <a href="https://huggingface.co/spaces/akhaliq/anycoder" target="_blank" class="attribution-link">
-                Built with anycoder
-            </a>
-        </div>
-    """)
-    
-    # Input Section
-    with gr.Column(elem_classes="input-section"):
-        prompt = gr.Textbox(
-            placeholder="Describe the image you want to create...",
-            lines=3,
-            max_lines=6,
-            label="",
-            show_label=False,
-            container=False,
-        )
-        
-        generate_btn = gr.Button(
-            "Generate",
-            variant="primary",
-            size="lg",
-            elem_classes="primary"
-        )
-    
-    # Output Section
-    with gr.Column(elem_classes="output-section"):
-        output_image = gr.Image(
-            type="pil",
-            label="",
-            show_label=False,
-            container=False,
-            buttons=["download"],
-        )
-    
-    # Footer
-    gr.HTML("""
-        <div class="footer-text">
-            <p>Powered by Z-Image Turbo from Tongyi-MAI</p>
-        </div>
-    """)
-    
-    # Event handlers
-    generate_btn.click(
-        fn=generate_image,
-        inputs=prompt,
-        outputs=output_image,
-        api_visibility="public"
-    )
-    
-    prompt.submit(
-        fn=generate_image,
-        inputs=prompt,
-        outputs=output_image,
-        api_visibility="public"
-    )
+    generate_btn.click(fn=generate_image, inputs=prompt_input, outputs=output_image)
+    prompt_input.submit(fn=generate_image, inputs=prompt_input, outputs=output_image)
 
 if __name__ == "__main__":
-    demo.launch(
-        share=False,
-        show_error=True,
-        theme=gr.themes.Soft(
-            primary_hue=gr.themes.colors.blue,
-            secondary_hue=gr.themes.colors.slate,
-            neutral_hue=gr.themes.colors.gray,
-            spacing_size=gr.themes.sizes.spacing_lg,
-            radius_size=gr.themes.sizes.radius_lg,
-            text_size=gr.themes.sizes.text_md,
-            font=[gr.themes.GoogleFont("Inter"), "SF Pro Display", "-apple-system", "BlinkMacSystemFont", "system-ui", "sans-serif"],
-            font_mono=[gr.themes.GoogleFont("JetBrains Mono"), "SF Mono", "ui-monospace", "monospace"],
-        ).set(
-            body_background_fill='#f5f5f7',
-            body_background_fill_dark='#000000',
-            button_primary_background_fill='#0071e3',
-            button_primary_background_fill_hover='#0077ed',
-            button_primary_text_color='#ffffff',
-            block_background_fill='#ffffff',
-            block_background_fill_dark='#1d1d1f',
-            block_border_width='0px',
-            block_shadow='0 2px 12px rgba(0, 0, 0, 0.08)',
-            block_shadow_dark='0 2px 12px rgba(0, 0, 0, 0.4)',
-            input_background_fill='#ffffff',
-            input_background_fill_dark='#1d1d1f',
-            input_border_width='1px',
-            input_border_color='#d2d2d7',
-            input_border_color_dark='#424245',
-            input_shadow='none',
-            input_shadow_focus='0 0 0 4px rgba(0, 113, 227, 0.15)',
-        ),
-        css=apple_css,
-    )
+    demo.launch(share=True)
+