import time
import torch
import numpy as np
import gradio as gr
from PIL import Image
from diffusers import SanaSprintPipeline

# ==========================================
# 1. SETUP & OPTIMIZATIONS
# ==========================================
device = "cpu"
torch.set_num_threads(torch.get_num_threads()) # Maximize CPU cores

print("Loading SANA-Sprint 0.6B Pipeline in 16-bit (bfloat16)...")
pipeline = SanaSprintPipeline.from_pretrained(
    "Efficient-Large-Model/Sana_Sprint_0.6B_1024px_diffusers",
    torch_dtype=torch.bfloat16
)
pipeline.to(device)
print("Pipeline Ready!")

# ==========================================
# 2. GENERATION FUNCTION
# ==========================================
def generate_image(prompt, seed):
    if not prompt:
        return None, "Please enter a prompt."
        
    start_time = time.time()
    
    # 2-Step Generation for Sana-Sprint
    image = pipeline(
        prompt=prompt,
        num_inference_steps=2, 
        guidance_scale=1.0,
        generator=torch.Generator(device=device).manual_seed(int(seed))
    ).images[0]
    
    elapsed_time = f"✅ Generated in {time.time() - start_time:.2f} seconds"
    return image, elapsed_time

# ==========================================
# 3. EDITING FUNCTION
# ==========================================
def edit_image(prompt, init_image, strength, seed):
    if not prompt:
        return None, "Please enter a prompt."
    if init_image is None:
        return None, "Please upload an image to edit."

    start_time = time.time()
    
    # Resize input image to 1024x1024 (SANA's native size)
    init_image = init_image.convert("RGB").resize((1024, 1024))
    
    # Convert image to Tensor
    image_tensor = torch.from_numpy(np.array(init_image)).float() / 127.5 - 1.0
    image_tensor = image_tensor.permute(2, 0, 1).unsqueeze(0).to(device)
    image_tensor = image_tensor.to(torch.bfloat16)
    
    with torch.no_grad():
        # Encode with SANA DC-AE
        latents = pipeline.vae.encode(image_tensor).latents
        latents = latents * pipeline.vae.config.scaling_factor
        
    # Generate noise and blend based on strength
    generator = torch.Generator(device=device).manual_seed(int(seed))
    noise = torch.randn_like(latents, generator=generator, dtype=torch.bfloat16)
    blended_latents = (1 - strength) * latents + strength * noise

    # Pass the blended latents into the pipeline
    edited_image = pipeline(
        prompt=prompt,
        num_inference_steps=2, 
        latents=blended_latents,
        guidance_scale=1.0,
        generator=generator
    ).images[0]
    
    elapsed_time = f"✅ Edited in {time.time() - start_time:.2f} seconds"
    return edited_image, elapsed_time

# ==========================================
# 4. GRADIO WEB UI
# ==========================================
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🚀 Fast SANA 0.6B on CPU")
    gr.Markdown("Ultra-fast Image Generation and Editing using **Sana-Sprint 0.6B** (2-step generation) in natively optimized 16-bit format.")

    with gr.Tabs():
        # TAB 1: TEXT TO IMAGE
        with gr.TabItem("🖼️ Generate Image"):
            with gr.Row():
                with gr.Column():
                    gen_prompt = gr.Textbox(label="Prompt", placeholder="A highly detailed cyberpunk city...")
                    gen_seed = gr.Slider(minimum=0, maximum=10000, step=1, value=42, label="Random Seed")
                    gen_button = gr.Button("Generate", variant="primary")
                with gr.Column():
                    gen_output = gr.Image(label="Generated Image")
                    gen_time = gr.Textbox(label="Generation Time", interactive=False)
            
            gen_button.click(
                fn=generate_image,
                inputs=[gen_prompt, gen_seed],
                outputs=[gen_output, gen_time]
            )

        # TAB 2: IMAGE TO IMAGE
        with gr.TabItem("🎨 Edit Image"):
            with gr.Row():
                with gr.Column():
                    edit_input_img = gr.Image(label="Upload Image to Edit", type="pil")
                    edit_prompt = gr.Textbox(label="Edit Prompt", placeholder="Change the background to winter...")
                    edit_strength = gr.Slider(minimum=0.1, maximum=1.0, step=0.05, value=0.6, label="Editing Strength (Higher = More changes)")
                    edit_seed = gr.Slider(minimum=0, maximum=10000, step=1, value=42, label="Random Seed")
                    edit_button = gr.Button("Edit Image", variant="primary")
                with gr.Column():
                    edit_output = gr.Image(label="Edited Image")
                    edit_time = gr.Textbox(label="Editing Time", interactive=False)
                    
            edit_button.click(
                fn=edit_image,
                inputs=[edit_prompt, edit_input_img, edit_strength, edit_seed],
                outputs=[edit_output, edit_time]
            )

# Launch the Gradio App
if __name__ == "__main__":
    demo.launch()