import torch
import spaces
import gradio as gr
from diffusers import ZImagePipeline, ZImageTransformer2DModel
from huggingface_hub import hf_hub_download

print("🔄 Initializing Z-Image-Turbo Pipeline...")

MODEL_FILENAME = "zimageTurboByStable_2603Fp8.safetensors"
REPO_ID = "ajsbsd/ZIT"

print(f"⬇️ Downloading {MODEL_FILENAME} from {REPO_ID}...")
model_path = hf_hub_download(repo_id=REPO_ID, filename=MODEL_FILENAME)

try:
    print("⚙️ Attempting to load as a FULL checkpoint (includes Text Encoders + VAE)...")
    # Load in bfloat16 to safely upcast FP8 weights and avoid missing CUDA kernels
    pipe = ZImagePipeline.from_single_file(
        model_path,
        torch_dtype=torch.bfloat16,
        low_cpu_mem_usage=True,
    )
    print("✅ Successfully loaded as full checkpoint!")
except Exception as e:
    print(f"⚠️ Full checkpoint load failed ({str(e)[:100]}...).")
    print("⚙️ Falling back to loading as TRANSFORMER ONLY...")
    
    pipe = ZImagePipeline.from_pretrained(
        "Tongyi-MAI/Z-Image-Turbo",
        torch_dtype=torch.bfloat16,
        low_cpu_mem_usage=True,
    )
    
    transformer = ZImageTransformer2DModel.from_single_file(
        model_path,
        torch_dtype=torch.bfloat16,
        low_cpu_mem_usage=True,
    )
    pipe.transformer = transformer
    print("✅ Successfully loaded custom transformer!")

# CRITICAL for ZeroGPU: Prevent Out-Of-Memory errors
pipe.enable_model_cpu_offload()
pipe.enable_attention_slicing()

print("🚀 Pipeline loaded and optimized for ZeroGPU! Ready to generate.")

@spaces.GPU
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))
    
    # Recommended settings: CFG 1.0 for Z-Image Turbo
    image = pipe(
        prompt=prompt,
        height=int(height),
        width=int(width),
        num_inference_steps=int(num_inference_steps),
        guidance_scale=1.0, 
        generator=generator,
    ).images[0]

    return image, int(seed)

examples = [
    ["Portrait of a young woman with natural skin texture, soft believable lighting, candid editorial style, highly detailed, photorealistic"],
    ["A candid full-body shot of a person walking in a softly lit urban street at dusk, natural real-life look, crisp faces, reliable anatomy"],
    ["Close-up portrait, natural real-photo realism, soft lighting, clean skin texture, no over-processed studio look, 85mm lens"]
]

custom_theme = gr.themes.Soft(primary_hue="emerald", secondary_hue="teal", neutral_hue="slate")

# Gradio 6.0+ Fix: Removed theme from Blocks constructor
with gr.Blocks(title="2603 ZIT — By Stable Yogi") as demo:
    gr.Markdown(
        """
        # 📷 2603 ZIT — By Stable Yogi
        **Fast photoreal Z-Image Turbo** with a natural, real-life look. Believable skin, faces, and lighting.
        
        **Recommended Settings:** Steps 8–9 · CFG 1.0 · Resolution 1152×896 or ~1024 square
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            prompt = gr.Textbox(label="✨ Prompt", placeholder="Describe the image you want to create...", lines=4)
            with gr.Accordion("⚙️ Advanced Settings", open=False):
                height = gr.Slider(minimum=512, maximum=1536, value=1152, step=64, label="Height")
                width = gr.Slider(minimum=512, maximum=1536, value=896, step=64, label="Width")
                num_inference_steps = gr.Slider(minimum=1, maximum=20, value=8, step=1, label="Inference Steps")
                with gr.Row():
                    randomize_seed = gr.Checkbox(label="🎲 Randomize Seed", value=True)
                    seed = gr.Number(label="Seed", value=42, precision=0, visible=False)
                randomize_seed.change(lambda r: gr.Number(visible=not r), inputs=[randomize_seed], outputs=[seed])
            generate_btn = gr.Button("🚀 Generate Image", variant="primary", size="lg")
            gr.Examples(examples=examples, inputs=[prompt], label="💡 Try these prompts")
        with gr.Column(scale=1):
            output_image = gr.Image(label="Generated Image", type="pil", format="png", height=600)
            used_seed = gr.Number(label="🎲 Seed Used", interactive=False)

    generate_btn.click(fn=generate_image, inputs=[prompt, height, width, num_inference_steps, seed, randomize_seed], outputs=[output_image, used_seed])
    prompt.submit(fn=generate_image, inputs=[prompt, height, width, num_inference_steps, seed, randomize_seed], outputs=[output_image, used_seed])

if __name__ == "__main__":
    # Gradio 6.0+ Fix: Pass theme to launch()
    demo.launch(theme=custom_theme)