import torch
import os
import psutil
import time
import math
import gradio as gr
from PIL import Image
from diffusers import AutoPipelineForImage2Image, AutoPipelineForText2Image

# --- System Resource Logic ---
def get_system_usage():
    cpu = psutil.cpu_percent(interval=None)
    ram = psutil.virtual_memory().percent
    # Returns values for the progress bars
    return cpu / 100, ram / 100, f"CPU: {cpu}% | RAM: {ram}%"

# --- Model Loading (Optimized for 16GB) ---
device = "cpu"
torch_dtype = torch.float32

pipe_kwargs = {"torch_dtype": torch_dtype, "use_safetensors": True}
i2i_pipe = AutoPipelineForImage2Image.from_pretrained("stabilityai/sdxl-turbo", **pipe_kwargs)
t2i_pipe = AutoPipelineForText2Image.from_pretrained("stabilityai/sdxl-turbo", **pipe_kwargs)

i2i_pipe.enable_attention_slicing()
t2i_pipe.enable_attention_slicing()
i2i_pipe.enable_vae_tiling() 
t2i_pipe.enable_vae_tiling()

def predict(init_image, prompt, strength, steps, seed):
    generator = torch.manual_seed(seed)
    if init_image is not None:
        # Simple square crop for better results
        w, h = init_image.size
        s = min(w, h)
        init_image = init_image.crop(((w-s)//2, (h-s)//2, (w+s)//2, (h+s)//2)).resize((512, 512))
        return i2i_pipe(prompt=prompt, image=init_image, generator=generator, num_inference_steps=int(steps), 
                        guidance_scale=0.0, strength=strength).images[0]
    else:
        return t2i_pipe(prompt=prompt, generator=generator, num_inference_steps=int(steps), 
                        guidance_scale=0.0, width=512, height=512).images[0]

# --- UI Construction ---
with gr.Blocks(css="#container{ max-width: 60rem; margin: 0 auto; }") as demo:
    with gr.Column(elem_id="container"):
        gr.Markdown("## 🚀 SDXL Turbo CPU + System Monitor")
        
        # System Usage Dashboard
        with gr.Row():
            with gr.Column():
                cpu_bar = gr.Label(label="System Status") # Text label for quick reading
            with gr.Column():
                cpu_plot = gr.Slider(label="CPU Load", interactive=False)
                ram_plot = gr.Slider(label="RAM Usage", interactive=False)

        with gr.Row():
            prompt = gr.Textbox(placeholder="A cinematic cat...", label="Prompt", scale=4)
            generate_bt = gr.Button("Generate", variant="primary", scale=1)
        
        with gr.Row():
            image_input = gr.Image(type="pil", label="Input (i2i)")
            image_output = gr.Image(label="Result")

        with gr.Accordion("Settings", open=False):
            strength = gr.Slider(0.0, 1.0, 0.7, label="Strength")
            steps = gr.Slider(1, 4, 2, step=1, label="Steps")
            seed = gr.Number(42, label="Seed")

        # Auto-update logic (Runs every 2 seconds)
        monitor_timer = gr.Timer(2)
        monitor_timer.tick(get_system_usage, outputs=[cpu_plot, ram_plot, cpu_bar])

        generate_bt.click(predict, [image_input, prompt, strength, steps, seed], image_output)

demo.launch()