from diffusers import StableDiffusionPipeline, DiffusionPipeline
import torch
import random
from datetime import datetime
import time
import gradio as gr
from PIL import Image

resolution = (512, 512)  # Risoluzione dell'immagine (width, height)
num_steps = 20
guidance_scale = 7.5
neg_prompt = "blurry"

model_id = "stablediffusionapi/duchaiten-real3d-nsfw-xl"
pipe = DiffusionPipeline.from_pretrained(model_id)

# Imposta il dispositivo su GPU se disponibile
device = "cuda" if torch.cuda.is_available() else "cpu"
pipe = pipe.to(device)


# Funzione per generare un'immagine
def generate_image(prompt, neg_prompt, width=resolution[0], height=resolution[1], seed=1, steps=35):
    generator = torch.manual_seed(seed)
    start_time = time.time()
    image = pipe(prompt, height=height, width=width, num_inference_steps=steps, guidance_scale=guidance_scale, generator=generator, negative_prompt=neg_prompt).images[0]
    end_time = time.time()    
    elapsed_time = end_time - start_time
    return image, elapsed_time

demo = gr.Interface(
    fn=generate_image,
    inputs=["text","text","number","number", "slider", "slider"],
    outputs=[gr.Image(),gr.Textbox(label="Elapsed time (seconds)")],
)

demo.launch()