import random
import gradio as gr
import requests
from concurrent.futures import ThreadPoolExecutor
from MonsterAPIClient import MClient
from typing import Tuple

client = MClient()

def generate_model_output(model: str, input_text: str, neg_prompt: str, samples: int, steps: int,
                          aspect_ratio: str, guidance_scale: float, random_seed: str) -> str:
    """
    Generate output from a specific model.

    Parameters:
        model (str): The name of the model.
        input_text (str): Your input text prompt.
        neg_prompt (str): Negative text prompt.
        samples (int): No. of images to be generated.
        steps (int): Sampling steps per image.
        aspect_ratio (str): Aspect ratio of the generated image.
        guidance_scale (float): Prompt guidance scale.
        random_seed (str): Random number used to initialize the image generation.

    Returns:
        str: The generated output text or image URL.
    """
    try:
        response = client.get_response(model, {
            "prompt": input_text,
            "negprompt": neg_prompt,
            "samples": samples,
            "steps": steps,
            "aspect_ratio": aspect_ratio,
            "guidance_scale": guidance_scale,
            "seed": random_seed,
        })
        output = client.wait_and_get_result(response['process_id'])
        if 'output' in output:
            return output['output']
        else:
            return "No output available."
    except Exception as e:
        return f"Error occurred: {str(e)}"

def generate_output(input_text: str, neg_prompt: str, samples: int, steps: int,
                    aspect_ratio: str, guidance_scale: float, random_seed: str):
    with ThreadPoolExecutor() as executor:
        # Schedule the function calls asynchronously
        future_sdxl_base = executor.submit(generate_model_output, 'sdxl-base', input_text, neg_prompt, samples, steps,
                                           aspect_ratio, guidance_scale, random_seed)
        future_txt2img = executor.submit(generate_model_output, 'txt2img', input_text, neg_prompt, samples, steps,
                                         aspect_ratio, guidance_scale, random_seed)
        
        # Get the results from the completed futures
        sdxl_base_output = future_sdxl_base.result()
        txt2img_output = future_txt2img.result()
    
    return [sdxl_base_output, txt2img_output]

# Function to stitch


input_components = [
    gr.inputs.Textbox(label="Input Prompt"),
    gr.inputs.Textbox(label="Negative Prompt"),
    gr.inputs.Slider(label="No. of Images to Generate", minimum=1, maximum=2, default=1, step = 1),
    gr.inputs.Slider(label="Sampling Steps per Image", minimum=30, maximum=40, default=30, step = 1),
    gr.inputs.Dropdown(label="Aspect Ratio", choices=["square", "landscape", "portrait"], default="square"),
    gr.inputs.Slider(label="Prompt Guidance Scale", minimum=0.1, maximum=20.0, default=7.5),
    gr.inputs.Textbox(label="Random Seed", default=random.randint(0, 1000000)),
]

output_component_sdxl_base = gr.Gallery(label="Stable Diffusion XL Output", type="pil", container = True)
output_component_txt2img = gr.Gallery(label="Stable Diffusion V1.5 Output", type="pil", container = True)

interface = gr.Interface(
    fn=generate_output,
    inputs=input_components,
    outputs=[output_component_sdxl_base, output_component_txt2img],
    live=False,
    capture_session=True,
    title="Stable Diffusion Evaluation [Beta] powered by MonsterAPI",
    description="""This HuggingFace Space lets you compare Stable-Diffusion V1.5 vs SDXL image quality. These models are hosted on [MonsterAPI](https://monsterapi.ai/?utm_source=sdxl-evaluation&utm_medium=referral) - An AI infrastructure platform for accessing LLMs via low-cost APIs and [no-code LLM finetuning](https://docs.monsterapi.ai/fine-tune-a-large-language-model-llm). MonsterAPI is powered by our low cost and highly scalable GPU computing platform - [Q Blocks](https://www.qblocks.cloud?utm_source=sdxl-evaluation&utm_medium=referral). Checkout our [SDXL API documentation](https://documenter.getpostman.com/view/13759598/2s8ZDVZ3Yi#37336bc8-7a4b-41fe-b253-6a6f7ba63c82) to get started.""",
    css="body {background-color: black}"
)

# Launch the Gradio app
interface.launch()