import torch
import tkinter as tk
from tkinter import filedialog, messagebox
from PIL import Image, ImageTk
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler

# Initialize Stable Diffusion
model_id = "runwayml/stable-diffusion-v1-5"
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.float16 if device == "cuda" else torch.float32

print(f"Using device: {device.upper()}")

pipeline = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=dtype)
pipeline.safety_checker = None  # WARNING: No safety filters—USE RESPONSIBLY
pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config)
pipeline = pipeline.to(device)
pipeline.enable_attention_slicing()

# GUI Setup
class RotonaGUI:
    def __init__(self, root):
        self.root = root
        self.root.title("Rotona - Unrestricted AI Image Generator")
        self.root.geometry("600x700")
        self.root.configure(bg="black")

        # Title Label
        self.title_label = tk.Label(root, text="Rotona AI Generator", fg="red", bg="black", font=("Helvetica", 18, "bold"))
        self.title_label.pack(pady=10)

        # Prompt Entry
        self.prompt_label = tk.Label(root, text="Enter your prompt:", fg="white", bg="black", font=("Helvetica", 12))
        self.prompt_label.pack()
        self.prompt_entry = tk.Entry(root, width=60)
        self.prompt_entry.pack(pady=5)

        # Generate Button
        self.generate_button = tk.Button(root, text="Generate Image", command=self.generate_image, bg="red", fg="white", font=("Helvetica", 12, "bold"))
        self.generate_button.pack(pady=10)

        # Canvas for Displaying Image
        self.canvas = tk.Canvas(root, width=512, height=512, bg="gray")
        self.canvas.pack(pady=10)

        # Save Button
        self.save_button = tk.Button(root, text="Save Image", command=self.save_image, bg="blue", fg="white", font=("Helvetica", 12, "bold"))
        self.save_button.pack(pady=10)
        self.save_button["state"] = "disabled"

        self.generated_image = None

    def generate_image(self):
        user_prompt = self.prompt_entry.get()
        if not user_prompt:
            messagebox.showerror("Error", "Please enter a prompt!")
            return

        try:
            negative_prompt = "blurry, distorted, bad anatomy, missing fingers, extra limbs, bad proportions, deformed hands, low quality"
            image = pipeline(user_prompt, negative_prompt=negative_prompt, guidance_scale=7.5, num_inference_steps=50).images[0]

            self.generated_image = image
            self.display_image(image)
            self.save_button["state"] = "normal"

        except torch.cuda.OutOfMemoryError:
            messagebox.showerror("Error", "Out of GPU memory! Try reducing image size.")
        except Exception as e:
            messagebox.showerror("Error", f"Image generation failed: {e}")

    def display_image(self, image):
        image.thumbnail((512, 512))
        img_tk = ImageTk.PhotoImage(image)
        self.canvas.create_image(256, 256, image=img_tk)
        self.canvas.image = img_tk

    def save_image(self):
        if self.generated_image:
            file_path = filedialog.asksaveasfilename(defaultextension=".png", filetypes=[("PNG files", "*.png")])
            if file_path:
                self.generated_image.save(file_path)
                messagebox.showinfo("Saved", f"Image saved successfully as {file_path}")

# Run GUI
root = tk.Tk()
app = RotonaGUI(root)
root.mainloop()