import gradio as gr
import os
import fal_client

def on_queue_update(update):
    if isinstance(update, fal_client.InProgress):
        for log in update.logs:
           print(log["message"])

def set_fal_key(api_key):
    os.environ["FAL_KEY"] = api_key
    return "FAL API key set successfully!"

def fal_recraft(prompt, api_key):
    set_fal_key(api_key=api_key)
    result = fal_client.subscribe(
        "fal-ai/recraft-v3",
        arguments={
            "prompt": prompt
        },
        with_logs=True,
        on_queue_update=on_queue_update,
    ) 
    return result['images'][0]['url']

css = """
#col-container {
    margin: 0 auto;
    max-width: 640px;
}
"""

examples = [
    ["this page contains the number of letters r that the word strawberry has",],
    ["this page contains the result of 2+5",],
    ["write 2 adjectives in english",],
    ["write the name of the US president",],
    ]


demo = gr.Interface(fal_recraft, 
                    inputs=[
                        gr.Textbox(label="Prompt"), 
                        gr.Textbox(type="password", label="FAL API Key", placeholder="You need to input your FAL API KEY here in order to access recraft-v3")
                        ], 
                    outputs="image", 
                    examples=examples,
                    title="Recraft V3 with FAL API and Gradio", 
                    description="""[Announcing red_panda by Recraft](https://blog.fal.ai/announcing-red_panda-by-recraft/)""",
                    cache_examples=False,
                    )

demo.launch()