app.py

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

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()
    
    image = pipeline(
        prompt=prompt,
        height=512,
        width=512,
        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 to 512x512
    init_image = init_image.convert("RGB").resize((512, 512))
    
    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():
        # FIX: Changed `.latents` to `.latent` (singular)
        latents = pipeline.vae.encode(image_tensor).latent
        latents = latents * pipeline.vae.config.scaling_factor
        
    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

    edited_image = pipeline(
        prompt=prompt,
        height=512,
        width=512,
        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 (512x512 Mode)")
    gr.Markdown("Running at 512x512 resolution to maximize speed on limited Cloud CPU cores.")

    with gr.Tabs():
        with gr.TabItem("🖼️ Generate Image"):
            with gr.Row():
                with gr.Column():
                    gen_prompt = gr.Textbox(label="Prompt", placeholder="A cyberpunk cat...")
                    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])

        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="Make it raining...")
                    edit_strength = gr.Slider(minimum=0.1, maximum=1.0, step=0.05, value=0.6, label="Editing Strength")
                    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])

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