File size: 3,529 Bytes
b1f81d8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 |
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()
|