import gradio as gr

import os
hf_token = os.environ.get("HF_TOKEN")

from gradio_client import Client

client = Client("https://fffiloni-gradio-api-tests.hf.space/", hf_token=hf_token)

examples = [
    "examples/image.jpg",
    "examples/audio.mp3", 
    "examples/video.mp4"
]

def api_t2i(prompt):
    result = client.predict(
		prompt,
		api_name="/text2image"
    )
    print(result)
    return result

def api_i2i(image):
    print(f"trying to send {image}")
    result = client.predict(
		image,
		api_name="/image2image"
    )
    print(result)
    return result

def api_t2a(prompt):
    result = client.predict(
		prompt,
		api_name="/text2audio"
    )
    print(result)
    return result

def api_a2a(audio):
    print(f"trying to send {audio}")
    result = client.predict(
		audio,
		api_name="/audio2audio"
    )
    print(result)
    return result

def api_t2v(prompt):
    result = client.predict(
		prompt,
		api_name="/text2video"
    )
    print(result)
    return result["video"]

def api_v2v(video):
    print(f"trying to send {video}")
    result = client.predict(
		video,
		api_name="/video2video"
    )
    print(result)
    return result["video"]

def api_t2f(prompt):
    result = client.predict(
		prompt,
		api_name="/text2files"
    )
    print(result)
    return result

def api_f2f(files):
    print(f"trying to send {files}")
    result = client.predict(
		files,
		api_name="/files2files"
    )
    print(result)
    return result

with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            prompt_in = gr.Textbox(label="Prompt IN", value="testing prompt")
            image_in = gr.Image(label="Image IN", type="filepath", value=examples[0])
            audio_in = gr.Audio(label="Audio IN", type="filepath", value=examples[1])
            video_in = gr.Video(label="Video IN", value=examples[2])
            files_in = gr.File(label="Files IN", value=examples, file_count="multiple")

            with gr.Row():
                prompt_t2i_btn = gr.Button("Run text 2 image")
                prompt_i2i_btn = gr.Button("Run image 2 image")
                prompt_t2a_btn = gr.Button("Run text 2 audio")
                prompt_a2a_btn = gr.Button("Run audio 2 audio")
                prompt_t2v_btn = gr.Button("Run text 2 video")
                prompt_v2v_btn = gr.Button("Run video 2 video")
                prompt_t2f_btn = gr.Button("Run text 2 files")
                prompt_f2f_btn = gr.Button("Run files 2 files")

        with gr.Column():
            image_out = gr.Image(label="Image OUT")
            audio_out = gr.Audio(label="Audio OUT")
            video_out = gr.Video(label="Video OUT")
            files_out = gr.File(label="Files OUT", file_count="multiple")

    prompt_t2i_btn.click(fn=api_t2i, inputs=[prompt_in], outputs=[image_out])
    prompt_i2i_btn.click(fn=api_i2i, inputs=[image_in], outputs=[image_out])
    prompt_t2a_btn.click(fn=api_t2a, inputs=[prompt_in], outputs=[audio_out])
    prompt_a2a_btn.click(fn=api_a2a, inputs=[audio_in], outputs=[audio_out])
    prompt_t2v_btn.click(fn=api_t2v, inputs=[prompt_in], outputs=[video_out])
    prompt_v2v_btn.click(fn=api_v2v, inputs=[video_in], outputs=[video_out])
    prompt_t2f_btn.click(fn=api_t2f, inputs=[prompt_in], outputs=[files_out])
    prompt_f2f_btn.click(fn=api_f2f, inputs=[files_in], outputs=[files_out])

demo.queue().launch()