import gradio as gr
from gradio_client import Client
import random

# Initialize the client for the Hugging Face Space
client = Client("ByteDance/Hyper-FLUX-8Steps-LoRA")

def generate_image(prompt, height, width, steps, scales, seed):
    """
    Generates an image based on the provided parameters by calling the Hugging Face Space API.

    Parameters:
    - prompt (str): The text prompt for image generation.
    - height (int): The height of the generated image.
    - width (int): The width of the generated image.
    - steps (int): The number of steps for the image generation process.
    - scales (float): The scaling factor.
    - seed (int): The seed for random number generation to ensure reproducibility.

    Returns:
    - result (str or Image): The generated image or a link to it.
    """
    # Generate a random seed if not provided
    if not seed:
        seed = random.randint(0, 100000)

    try:
        result = client.predict(
            height=int(height),
            width=int(width),
            steps=int(steps),
            scales=float(scales),
            prompt=prompt,
            seed=int(seed),
            api_name="/process_image"
        )
        return result
    except Exception as e:
        return f"An error occurred: {e}"

# Define the Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# Hyper-FLUX-8Steps-LoRA Image Generator")
    gr.Markdown("Generate images based on your text prompts using the Hyper-FLUX-8Steps-LoRA model.")

    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(
                label="Prompt",
                placeholder="Enter your descriptive text here...",
                lines=2
            )
            height = gr.Number(
                label="Height",
                value=1024,
                precision=0,
                interactive=True
            )
            width = gr.Number(
                label="Width",
                value=1024,
                precision=0,
                interactive=True
            )
            steps = gr.Number(
                label="Steps",
                value=8,
                precision=0,
                interactive=True
            )
            scales = gr.Number(
                label="Scale",
                value=3.5,
                precision=1,
                interactive=True
            )
            seed = gr.Number(
                label="Seed",
                value=3413,
                precision=0,
                interactive=True
            )
            generate_button = gr.Button("Generate Image")

        with gr.Column():
            output_image = gr.Image(label="Generated Image", interactive=False)

    # Define the button click action
    generate_button.click(
        fn=generate_image,
        inputs=[prompt, height, width, steps, scales, seed],
        outputs=output_image
    )

    # Optional: Add a footer or additional information
    gr.Markdown("© 2024 Your Name. All rights reserved.")

# Launch the Gradio app
if __name__ == "__main__":
    demo.launch()