import gradio as gr
import random

# Load each model separately
model1 = gr.load("models/pimpilikipilapi1/NSFW_master")
model2 = gr.load("models/prashanth970/flux-lora-uncensored")
model3 = gr.load("models/DiegoJR1973/NSFW-TrioHMH-Flux")

def generate_images(text, seed, width, height, guidance_scale, num_inference_steps):
    if seed is not None:
        random.seed(seed)

    # Generate images using each model with the text prompt only
    result_image1 = model1(text)
    result_image2 = model2(text)
    result_image3 = model3(text)

    # Print parameters for debugging
    print(f"Text: {text}, Seed: {seed}, Width: {width}, Height: {height}, Guidance Scale: {guidance_scale}, Steps: {num_inference_steps}")
    
    return result_image1, result_image2, result_image3

def randomize_parameters():
    seed = random.randint(0, 999999)
    width = random.randint(512, 2048)
    height = random.randint(512, 2048)
    guidance_scale = round(random.uniform(0.1, 20.0), 1)
    num_inference_steps = random.randint(1, 40)
    
    return seed, width, height, guidance_scale, num_inference_steps

interface = gr.Interface(
    fn=generate_images,
    inputs=[
        gr.Textbox(label="Type here your imagination:", placeholder="Type or click an example..."),
        gr.Slider(label="Seed", minimum=0, maximum=999999, step=1),
        gr.Slider(label="Width", minimum=512, maximum=2048, step=64, value=1024),
        gr.Slider(label="Height", minimum=512, maximum=2048, step=64, value=1024),
        gr.Slider(label="Guidance Scale", minimum=0.1, maximum=20.0, step=0.1, value=3.0),
        gr.Slider(label="Number of inference steps", minimum=1, maximum=40, step=1, value=28),
    ],
    outputs=[
        gr.Image(label="Generated Image from NSFW_master"),
        gr.Image(label="Generated Image from flux-lora-uncensored"),
        gr.Image(label="Generated Image from NSFW-TrioHMH-Flux")
    ],
    description="Generate images with three different models based on your prompt. Note: the models run on the CPU, which may affect performance.",
)

interface.launch()