import gradio as gr
import threading
import os
import torch

os.environ["OMP_NUM_THREADS"] = str(os.cpu_count())
torch.set_num_threads(os.cpu_count())

model1 = gr.load("models/prithivMLmods/SD3.5-Turbo-Realism-2.0-LoRA")
model2 = gr.load("models/Purz/face-projection")

stop_event = threading.Event()

def generate_images(text, selected_model, steps, cfg_scale, seed, width, height):
    stop_event.clear()

    if selected_model == "Model 1 (Turbo Realism)":
        model = model1
    elif selected_model == "Model 2 (Face Projection)":
        model = model2
    else:
        return ["Invalid model selection."] * 3

    # Convert seed to integer (handle empty/None)
    try:
        seed = int(seed) if seed not in [None, ""] else -1
    except:
        seed = -1

    results = []
    for i in range(3):
        if stop_event.is_set():
            return ["Image generation stopped by user."] * 3

        modified_text = f"{text} variation {i+1}"
        result = model(
            modified_text,
            #num_inference_steps=int(steps),
            guidance_scale=float(cfg_scale),
            height=int(height),
            width=int(width),
            seed=seed if seed != -1 else None
        )
        results.append(result)

    return results

def stop_generation():
    """Stops the ongoing image generation by setting the stop_event flag."""
    stop_event.set()
    return ["Generation stopped."] * 3

with gr.Blocks() as interface:
    gr.Markdown(
        "### ⚠ Sorry for the inconvenience. The Space is currently running on the CPU, which might affect performance. We appreciate your understanding."
    )
    
    with gr.Row():
        text_input = gr.Textbox(label="Prompt", placeholder="Type your imagination here...")
        model_selector = gr.Radio(
            ["Model 1 (Turbo Realism)", "Model 2 (Face Projection)"],
            label="Model Selection",
            value="Model 1 (Turbo Realism)"
        )
    
    with gr.Accordion("Advanced Parameters", open=False):
        with gr.Row():
            steps = gr.Slider(1, 150, value=25, label="Inference Steps", info="(20-50 for quality/speed balance)")
            cfg_scale = gr.Slider(1.0, 20.0, value=7.5, label="CFG Scale", info="(7-12 for good balance)")
            seed = gr.Number(label="Seed", value=-1, precision=0, info="-1 for random")
        
        with gr.Row():
            width = gr.Dropdown(
                choices=["512", "640", "768", "896", "1024"],
                value="512",
                label="Width",
                allow_custom_value=True
            )
            height = gr.Dropdown(
                choices=["512", "640", "768", "896", "1024"],
                value="512",
                label="Height",
                allow_custom_value=True
            )
    
    with gr.Row():
        generate_button = gr.Button("Generate 3 Images 🎨", variant="primary")
        stop_button = gr.Button("Stop Generation", variant="stop")
    
    with gr.Row():
        output1 = gr.Image(label="Variant 1", type="pil", show_label=True)
        output2 = gr.Image(label="Variant 2", type="pil", show_label=True)
        output3 = gr.Image(label="Variant 3", type="pil", show_label=True)

    generate_button.click(
        generate_images,
        inputs=[text_input, model_selector, steps, cfg_scale, seed, width, height],
        outputs=[output1, output2, output3]
    )
    stop_button.click(
        stop_generation,
        inputs=[],
        outputs=[output1, output2, output3]
    )

interface.launch()