import gradio as gr
from gradio_client import Client

# Function to call /update_inference_count API
def update_inference_count():
    client = Client("https://nv-sana.mit.edu/")
    result = client.predict(api_name="/update_inference_count")
    return result

# Function to call /run_inference API
def run_inference(num_imgs):
    client = Client("https://nv-sana.mit.edu/")
    result = client.predict(num_imgs=num_imgs, api_name="/run_inference")
    return result

# Function to call /run API
def run_image_generation(prompt, negative_prompt, style, use_negative_prompt, num_imgs, seed, height, width, flow_dpms_guidance_scale, flow_dpms_pag_guidance_scale, flow_dpms_inference_steps, randomize_seed):
    client = Client("https://nv-sana.mit.edu/")
    result = client.predict(
        prompt=prompt,
        negative_prompt=negative_prompt,
        styles=style,
        use_negative_prompt=use_negative_prompt,
        num_imgs=num_imgs,
        seed=seed,
        height=height,
        width=width,
        flow_dpms_guidance_scale=flow_dpms_guidance_scale,
        flow_dpms_pag_guidance_scale=flow_dpms_pag_guidance_scale,
        flow_dpms_inference_steps=flow_dpms_inference_steps,
        randomize_seed=randomize_seed,
        api_name="/run"
    )
    return result

# Create Gradio interface for /update_inference_count
update_iface = gr.Interface(
    fn=update_inference_count,
    inputs=None,
    outputs="text",
    title="Update Inference Count",
    description="Click the button to update the inference count."
)

# Create Gradio interface for /run_inference
run_inference_iface = gr.Interface(
    fn=run_inference,
    inputs=gr.Number(label="Number of Images", value=1),
    outputs="text",
    title="Run Inference",
    description="Run inference with the specified number of images."
)

# Create Gradio interface for /run
run_image_generation_iface = gr.Interface(
    fn=run_image_generation,
    inputs=[
        gr.Textbox(label="Prompt"),
        gr.Textbox(label="Negative Prompt"),
        gr.Radio(choices=["No style", "Cinematic", "Photographic", "Anime", "Manga", "Digital Art"], label="Image Style"),
        gr.Checkbox(label="Use Negative Prompt"),
        gr.Number(label="Number of Images", value=1),
        gr.Number(label="Seed", value=0),
        gr.Number(label="Height", value=1024),
        gr.Number(label="Width", value=1024),
        gr.Number(label="CFG Guidance Scale", value=5),
        gr.Number(label="PAG Guidance Scale", value=2),
        gr.Number(label="Sampling Steps", value=18),
        gr.Checkbox(label="Randomize Seed")
    ],
    outputs="text",
    title="Run Image Generation",
    description="Generate images with the specified parameters."
)

# Combine interfaces into a single app with tabs
demo = gr.TabbedInterface(
    [update_iface, run_inference_iface, run_image_generation_iface],
    ["Update Inference Count", "Run Inference", "Run Image Generation"]
)

# Run the Gradio app
demo.launch()