import gradio as gr
import threading
import os
import torch

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

# Load only the Turbo Realism model
model = gr.load("models/OminduAnjana/HarmonyVerse-3.2")
stop_event = threading.Event()

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

with gr.Blocks() as interface:
    # Warning message
    gr.Markdown(
        "### ⚠ Sorry for the inconvenience. The Space is currently running on the CPU, which might affect performance. We appreciate your understanding."
    )
    
    # Input area
    text_input = gr.Textbox(
        label="Type here your imagination:", 
        placeholder="Type your prompt..."
    )
    
    # Buttons
    with gr.Row():
        generate_button = gr.Button("Generate 3 Images 🎨")
        stop_button = gr.Button("Stop Image Generation")
    
    # Output images
    with gr.Row():
        output1 = gr.Image(label="Generated Image 1")
        output2 = gr.Image(label="Generated Image 2")
        output3 = gr.Image(label="Generated Image 3")
    
    # Set up the event handlers
    generate_button.click(
        generate_images, 
        inputs=[text_input], 
        outputs=[output1, output2, output3]
    )
    
    stop_button.click(
        stop_generation, 
        inputs=[], 
        outputs=[output1, output2, output3]
    )

interface.launch()