import torch
import gradio as gr
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler

# Load Model
device = "cuda" if torch.cuda.is_available() else "cpu"
model_id = "runwayml/stable-diffusion-v1-5"
pipeline = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16 if device == "cuda" else torch.float32)
pipeline.safety_checker = None
pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config)
pipeline = pipeline.to(device)

# Image generation function
def generate_image(prompt):
    negative_prompt = "blurry, distorted, bad anatomy, missing fingers, extra limbs, bad proportions, deformed hands, low quality, cropped, out of frame, partial body, cut off, head cut off"
    refined_prompt = f"{prompt}, centered, full body, well-framed, symmetrical"

    with torch.no_grad():
        image = pipeline(refined_prompt, negative_prompt=negative_prompt, guidance_scale=7.5, num_inference_steps=50).images[0]

    torch.cuda.empty_cache()  # Free GPU memory
    return image

# Gradio UI
interface = gr.Interface(
    fn=generate_image,
    inputs=gr.Textbox(label="Enter Your Prompt"),
    outputs=gr.Image(label="Generated Image"),
    title="Rotona: Unrestricted Image Generator",
    description="Enter a prompt to generate an AI image. This model is unrestricted—use responsibly."
)

# Launch
if __name__ == "__main__":
    interface.launch()