app.py

import spaces
import gradio as gr
import torch
from diffusers import ZImagePipeline, AutoPipelineForText2Image
import os
from pathlib import Path

# Load the base 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!")

# Store the current model state
current_model = "base"
lora_loaded = False

@spaces.GPU()
def generate_image(
    prompt,
    model_choice,
    progress=gr.Progress(track_tqdm=True)
):
    """
    Generate an image using Z-Image Turbo model with optional LoRA.
    
    Args:
        prompt: Text description of the desired image
        model_choice: Either "Base Model" or "Classic Painting LoRA"
    
    Returns:
        Generated PIL Image
    """
    global pipe, current_model, lora_loaded
    
    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"
    
    # Handle model switching
    progress(0.05, desc="Loading model...")
    
    try:
        if model_choice == "Classic Painting LoRA" and current_model != "lora":
            # Load LoRA weights
            if not lora_loaded:
                print("Loading Classic Painting LoRA...")
                pipe.load_lora_weights(
                    "renderartist/Classic-Painting-Z-Image-Turbo-LoRA",
                    adapter_name="classic_painting",
                    weight_name="pytorch_lora_weights.safetensors"
                )
                lora_loaded = True
            
            # Set LoRA adapter
            pipe.set_adapters(["classic_painting"], adapter_weights=[0.8])
            current_model = "lora"
            progress(0.15, desc="LoRA loaded, generating image...")
            
            # Generate with LoRA settings
            generator = torch.Generator(device).manual_seed(42)
            result = pipe(
                prompt=prompt,
                negative_prompt=None,
                height=1024,
                width=1024,
                num_inference_steps=9,
                guidance_scale=0.0,
                generator=generator,
            )
        
        elif model_choice == "Base Model" and current_model != "base":
            # Disable LoRA
            pipe.disable_lora()
            current_model = "base"
            progress(0.15, desc="Generating image...")
            
            # Generate with base model settings
            generator = torch.Generator(device).manual_seed(42)
            result = pipe(
                prompt=prompt,
                negative_prompt=None,
                height=1024,
                width=1024,
                num_inference_steps=9,
                guidance_scale=0.0,
                generator=generator,
            )
        
        else:
            # Model already loaded, just generate
            progress(0.15, desc="Generating image...")
            generator = torch.Generator(device).manual_seed(42)
            result = pipe(
                prompt=prompt,
                negative_prompt=None,
                height=1024,
                width=1024,
                num_inference_steps=9,
                guidance_scale=0.0,
                generator=generator,
            )
        
        image = result.images[0]
        progress(1.0, desc="Complete!")
        
        return image
    
    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);
}

/* Model Selector */
.model-selector {
    margin-bottom: 24px;
}

.model-selector label {
    font-size: 15px !important;
    font-weight: 500 !important;
    color: #1d1d1f !important;
    margin-bottom: 8px !important;
    display: block !important;
}

/* 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;
}

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;
}

.dark .model-selector label {
    color: #f5f5f7 !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"):
        # Model Selector
        with gr.Group(elem_classes="model-selector"):
            model_choice = gr.Radio(
                choices=["Base Model", "Classic Painting LoRA"],
                value="Base Model",
                label="Select Model",
                info="Choose between the base Z-Image Turbo model or the Classic Painting LoRA variant"
            )
        
        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>
            <p>Classic Painting LoRA by renderartist</p>
        </div>
    """)
    
    # Event handlers
    generate_btn.click(
        fn=generate_image,
        inputs=[prompt, model_choice],
        outputs=output_image,
        api_visibility="public"
    )
    
    prompt.submit(
        fn=generate_image,
        inputs=[prompt, model_choice],
        outputs=output_image,
        api_visibility="public"
    )

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,
    )