import torch
import gradio as gr
import random
from diffusers import DiffusionPipeline
from sdnq import SDNQConfig
from sdnq.loader import apply_sdnq_options_to_model

# --- Model Configuration ---
MODEL_ID = "Abrahamm3r/Z-Image-SDNQ-uint4-svd-r32"
device = "cuda" if torch.cuda.is_available() else "cpu"

print(f"Loading model: {MODEL_ID} on {device}...")

# 1. Load the pipeline with trust_remote_code=True for Z-Image architecture
# We use bfloat16 as it is standard for these newer flux/z-image models
pipe = DiffusionPipeline.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.bfloat16,
    trust_remote_code=True
)

# 2. Apply SDNQ quantization hooks to the transformer
# This is critical for the model to run with the compressed weights
pipe.transformer = apply_sdnq_options_to_model(pipe.transformer)

# 3. Optimize memory
if device == "cuda":
    pipe.to(device)
    # Enable if you are on a smaller GPU (e.g., T4 16GB) to save VRAM
    # pipe.enable_model_cpu_offload() 
    
print("Model loaded successfully!")

# --- Helper Functions ---

# Preset resolutions for Aspect Ratios
# Z-Image handles various resolutions, but these are safe standard presets
ASPECT_RATIOS = {
    "1:1 (Square)": (1024, 1024),
    "16:9 (Cinematic)": (1280, 720),
    "9:16 (Portrait)": (720, 1280),
    "4:3 (Photo)": (1152, 864),
    "3:4 (Portrait Photo)": (864, 1152),
    "21:9 (Ultrawide)": (1536, 640)
}

def generate_image(prompt, negative_prompt, steps, aspect_ratio_choice, seed, guidance_scale):
    # Determine Width/Height from preset
    width, height = ASPECT_RATIOS.get(aspect_ratio_choice, (1024, 1024))
    
    # Handle Seed
    if seed == -1:
        seed = random.randint(0, 2**32 - 1)
    generator = torch.Generator(device=device).manual_seed(int(seed))
    
    print(f"Generating: '{prompt}' | Steps: {steps} | Size: {width}x{height} | Seed: {seed}")
    
    try:
        image = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            width=width,
            height=height,
            num_inference_steps=steps,
            guidance_scale=guidance_scale,
            generator=generator
        ).images[0]
        return image, seed
    except Exception as e:
        raise gr.Error(f"Generation failed: {str(e)}")

# --- Gradio UI ---

custom_css = """
#col-container { max-width: 800px; margin: 0 auto; }
#generate-btn { background: linear-gradient(90deg, #4B79A1 0%, #283E51 100%); border: none; color: white; }
"""

with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(f"# ⚡ Z-Image SDNQ (uint4-svd-r32) Generator")
        gr.Markdown(f"Running `{MODEL_ID}`. This uses Structured Decomposable Neural Quantization for high efficiency.")
        
        with gr.Row():
            with gr.Column(scale=2):
                prompt = gr.Textbox(label="Prompt", placeholder="Describe the image you want...", lines=3)
                negative_prompt = gr.Textbox(label="Negative Prompt", placeholder="Low quality, blurry, ugly...", value="low quality, bad anatomy, worst quality, distortion, blurry")
                
                with gr.Row():
                    aspect_ratio = gr.Dropdown(
                        label="Aspect Ratio", 
                        choices=list(ASPECT_RATIOS.keys()), 
                        value="1:1 (Square)"
                    )
                    steps = gr.Slider(label="Inference Steps", minimum=4, maximum=50, step=1, value=25)
                
                with gr.Row():
                    guidance = gr.Slider(label="Guidance Scale", minimum=1.0, maximum=10.0, step=0.1, value=3.5)
                    seed = gr.Number(label="Seed (-1 for Random)", value=-1, precision=0)
                
                run_btn = gr.Button("Generate Image", elem_id="generate-btn", size="lg")
            
            with gr.Column(scale=2):
                result_image = gr.Image(label="Generated Image", type="pil")
                seed_output = gr.Label(label="Used Seed")

    run_btn.click(
        fn=generate_image,
        inputs=[prompt, negative_prompt, steps, aspect_ratio, seed, guidance],
        outputs=[result_image, seed_output]
    )

if __name__ == "__main__":
    demo.queue().launch()