import gradio as gr
import random


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)

    result_image1 = model1(
        text, 
        width=width, 
        height=height, 
        guidance_scale=guidance_scale, 
        num_inference_steps=num_inference_steps
    )

    result_image2 = model2(
        text, 
        width=width - 128 if width > 640 else width,
        height=height - 128 if height > 640 else height,
        guidance_scale=guidance_scale * 1.2,
        num_inference_steps=max(1, num_inference_steps - 5)
    )
    

    result_image3 = model3(
        text, 
        width=width + 128 if width < 1920 else width,
        height=height + 128 if height < 1920 else height,
        guidance_scale=max(0.1, guidance_scale * 0.8),
        num_inference_steps=min(40, num_inference_steps + 5)
    )

    print(f"Model 1: Width={width}, Height={height}, Guidance Scale={guidance_scale}, Steps={num_inference_steps}")
    print(f"Model 2: Width={width - 128}, Height={height - 128}, Guidance Scale={guidance_scale * 1.2}, Steps={max(1, num_inference_steps - 5)}")
    print(f"Model 3: Width={width + 128}, Height={height + 128}, Guidance Scale={max(0.1, guidance_scale * 0.8)}, Steps={min(40, num_inference_steps + 5)}")
    
    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 01"),
        gr.Image(label="Generated Image 02"),
        gr.Image(label="Generated Image 03")
    ],
    description="Generate images with three different models, each with slight variations. Please note that the models are running on the CPU, which might affect performance. Thank you for your patience!",
)

interface.launch()