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):
    stop_event.clear()

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

    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)
        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

css = """
.container {
    max-width: 1200px;
    margin: auto;
    padding: 20px;
}
.main-title {
    text-align: center;
    color: #1a73e8;
    font-size: 2.5em;
    margin-bottom: 1em;
    font-weight: 700;
}
.notice {
    background-color: #fef3c7;
    border-left: 4px solid #f59e0b;
    padding: 1em;
    margin: 1em 0;
    border-radius: 4px;
}
.input-group {
    background-color: #f8fafc;
    padding: 2em;
    border-radius: 12px;
    box-shadow: 0 4px 6px -1px rgba(0,0,0,0.1);
    margin: 1em 0;
}
.button-primary {
    background-color: #1a73e8 !important;
}
.button-secondary {
    background-color: #dc2626 !important;
}
.output-images {
    gap: 1em;
    margin-top: 2em;
}
.output-image {
    border-radius: 8px;
    overflow: hidden;
    box-shadow: 0 4px 6px -1px rgba(0,0,0,0.1);
}
"""

with gr.Blocks(css=css, theme=gr.themes.Soft()) as interface:
    with gr.Column(elem_classes="container"):
        gr.Markdown("# AI Image Generator 🎨", elem_classes="main-title")
        
        gr.Markdown(
            "⚠️ Note: Currently running on CPU mode which may affect performance.",
            elem_classes="notice"
        )
        
        with gr.Group(elem_classes="input-group"):
            text_input = gr.Textbox(
                label="Describe your imagination",
                placeholder="Enter a detailed description of the image you want to generate...",
                lines=3
            )
            model_selector = gr.Radio(
                ["Turbo Realism", "Face Projection"],
                label="Select Model",
                value="Turbo Realism",
                container=False
            )
            
            with gr.Row():
                generate_button = gr.Button(
                    "✨ Generate Images",
                    elem_classes="button-primary"
                )
                stop_button = gr.Button(
                    "🛑 Stop Generation",
                    elem_classes="button-secondary"
                )
        
        with gr.Row(elem_classes="output-images"):
            output1 = gr.Image(label="Variation 1", elem_classes="output-image")
            output2 = gr.Image(label="Variation 2", elem_classes="output-image")
            output3 = gr.Image(label="Variation 3", elem_classes="output-image")
        
        gr.Markdown(
            """
            ### Tips for better results:
            - Be specific in your descriptions
            - Include details about style, mood, and lighting
            - Experiment with different models for varied results
            """,
        )
    
    generate_button.click(
        generate_images,
        inputs=[text_input, model_selector],
        outputs=[output1, output2, output3]
    )
    stop_button.click(
        stop_generation,
        inputs=[],
        outputs=[output1, output2, output3]
    )

interface.launch()