import tkinter as tk
from tkinter import scrolledtext, messagebox
import torch
from diffusers import DiffusionPipeline
import subprocess
import sys

# Function to install packages needed for the environment
def install_packages():
    try:
        subprocess.check_call([sys.executable, "-m", "pip", "install", "transformers", "accelerate", "safetensors"])
        messagebox.showinfo("Success", "Packages installed successfully!")
    except subprocess.CalledProcessError:
        messagebox.showerror("Error", "Failed to install packages.")

# Function to check and return appropriate device (GPU or iGPU)
def get_device():
    # Check if CUDA is available (for MX or other CUDA-enabled NVIDIA GPUs)
    if torch.cuda.is_available():
        # Get the name of the available GPU
        device_name = torch.cuda.get_device_name(0).lower()
        if "mx" in device_name:  # Prioritize MX GPUs if detected
            print(f"Using NVIDIA MX GPU: {device_name}")
            return "cuda"
        else:
            # If an NVIDIA GPU but not MX (could be a more powerful GPU), use it
            print(f"Using CUDA-enabled GPU: {device_name}")
            return "cuda"
    
    # If no CUDA GPU, check for Integrated GPU (iGPU) support
    # Integrated GPUs won't support CUDA, but we can still run inference without it
    else:
        print("No CUDA device detected. Trying to run with iGPU (integrated GPU).")
        return "cpu"  # Use CPU fallback for iGPU or lack of CUDA support

# Initialize the device (GPU or iGPU)
device = get_device()

# Load the diffusion pipeline and optimize it for the target device
pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", use_safetensors=True, variant="fp16")
pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
pipe.enable_model_cpu_offload()

# Load both base & refiner models to the target device (GPU or iGPU)
base = DiffusionPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", variant="fp16", use_safetensors=True
)
base.to(device)

refiner = DiffusionPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-refiner-1.0",
    text_encoder_2=base.text_encoder_2,
    vae=base.vae,
    variant="fp16",
    use_safetensors=True,
)
refiner.to(device)

# Set parameters for inference
n_steps = 40
high_noise_frac = 0.8

def run_inference():
    main_prompt = main_prompt_text.get("1.0", tk.END).strip()
    negative_prompt = negative_prompt_text.get("1.0", tk.END).strip() if negative_prompt_text.get("1.0", tk.END).strip() else None
    
    # Run the pipeline with the prompts
    result = pipe(main_prompt, negative_prompt=negative_prompt)
    output_text.delete("1.0", tk.END)
    output_text.insert(tk.END, str(result))

    # Run both experts (base and refiner)
    image = base(
        prompt=main_prompt,
        num_inference_steps=n_steps,
        denoising_end=high_noise_frac,
        output_type="latent",
    ).images
    image = refiner(
        prompt=main_prompt,
        num_inference_steps=n_steps,
        denoising_start=high_noise_frac,
        image=image,
    ).images[0]

    # Display the image or result as needed
    output_text.insert(tk.END, "Image generated successfully.")

# Create the main window
root = tk.Tk()
root.title("Inference GUI")

# Main prompt input
tk.Label(root, text="Main Prompt:").pack()
main_prompt_text = scrolledtext.ScrolledText(root, wrap=tk.WORD, width=50, height=10)
main_prompt_text.pack()

# Negative prompt input
tk.Label(root, text="Negative Prompt (optional):").pack()
negative_prompt_text = scrolledtext.ScrolledText(root, wrap=tk.WORD, width=50, height=5)
negative_prompt_text.pack()

# Run button
run_button = tk.Button(root, text="Run Inference", command=run_inference)
run_button.pack()

# Install button
install_button = tk.Button(root, text="Install Packages", command=install_packages)
install_button.pack()

# Output text area
output_text = scrolledtext.ScrolledText(root, wrap=tk.WORD, width=50, height=10)
output_text.pack()

# Start the GUI event loop
root.mainloop()