import gradio as gr
import jax

print(f"JAX devices: {jax.devices()}")
print(f"JAX device type: {jax.devices()[0].device_kind}")

def infer_charred(prompt):
    # your inference function for charr stable difusion control
    return None


with gr.Blocks(theme="gradio/soft") as demo:

    gr.Markdown("# Character-Aware Stable Diffusion (CHARRED)")

    with gr.Tab("Journal"):
        gr.Markdown(
            "## On How Four Crazy Fellows Embarked on Training a U-Net from Scratch in Five Days and Almost Died in the End."
        )

    with gr.Tab("☢️ DEMO ☢️"):
        gr.Markdown(
            "## This is a demo of the CHARRED character-aware stable diffusion model for you to enjoy at your own leisure, risk and peril"
        )
        prompt_input_charr = gr.Textbox(label="Prompt")
        charred_output = gr.Image(label="Output Image")
        submit_btn = gr.Button(value="Submit")
        charred_inputs = [prompt_input_charr]
        submit_btn.click(
            fn=infer_charred, inputs=charred_inputs, outputs=[charred_output]
        )
    # examples = [["postage stamp from california", "low quality", "charr_output.png", "charr_output.png" ]]
    # gr.Examples(fn = infer_sd, inputs = ["text", "text", "image", "image"], examples=examples, cache_examples=True)


demo.launch()