import gradio as gr

import torch

from diffusers import UniDiffuserPipeline


device = 'cuda' if torch.cuda.is_available() else 'cpu'
model_id = "thu-ml/unidiffuser-v1"
# model_id = "thu-ml/unidiffuser-v0"
pipeline = UniDiffuserPipeline.from_pretrained(
    model_id,
)
pipeline.to(device)


def convert_to_none(s):
    if s:
        return s
    else:
        return None


def set_mode(mode):
    if mode == "joint":
        pipeline.set_joint_mode()
    elif mode == "text2img":
        pipeline.set_text_to_image_mode()
    elif mode == "img2text":
        pipeline.set_image_text_mode()
    elif mode == "text":
        pipeline.set_text_mode()
    elif mode == "img":
        pipeline.set_image_mode()


def sample(mode, prompt, image, num_inference_steps, guidance_scale, seed):
    set_mode(mode)
    prompt = convert_to_none(prompt)
    image = convert_to_none(image)
    generator = torch.Generator(device=device).manual_seed(seed)
    output_sample = pipeline(
        prompt=prompt,
        image=image,
        num_inference_steps=num_inference_steps,
        guidance_scale=guidance_scale,
        generator=generator,
    )
    sample_image = None
    sample_text = ""
    if output_sample.images is not None:
        sample_image = output_sample.images[0]
    if output_sample.text is not None:
        sample_text = output_sample.text[0]
    return sample_image, sample_text
    

iface = gr.Interface(
    fn=sample,
    inputs=[
        gr.Textbox(value="", label="Generation Task"),
        gr.Textbox(value="", label="Conditioning prompt"),
        gr.Image(value=None, label="Conditioning image", type="pil"),
        gr.Number(value=20, label="Num Inference Steps", precision=0),
        gr.Number(value=8.0, label="Guidance Scale"),
        gr.Number(value=0, label="Seed", precision=0),
    ],
    outputs=[
        gr.Image(label="Sample image"),
        gr.Textbox(label="Sample text"),
    ],
)
iface.launch()

# from unidiffuser.sample_v0 import sample
# from unidiffuser.sample_v0_test import sample
# from unidiffuser.sample_v1 import sample
# from unidiffuser.sample_v1_test import sample


# def predict(mode, prompt, image, sample_steps, guidance_scale, seed):
#     output_images, output_text = sample(
#         mode, prompt, image, sample_steps=sample_steps, scale=guidance_scale, seed=seed,
#     )
#     sample_image = None
#     sample_text = ""
#     if output_images is not None:
#         sample_image = output_images[0]
#     if output_text is not None:
#         sample_text = output_text[0]
#     return sample_image, sample_text


# iface = gr.Interface(
#     fn=predict,
#     inputs=[
#         gr.Textbox(value="", label="Generation Task"),
#         gr.Textbox(value="", label="Conditioning prompt"),
#         gr.Image(value=None, label="Conditioning image", type="filepath"),
#         gr.Number(value=50, label="Num Inference Steps", precision=0),
#         gr.Number(value=7.0, label="Guidance Scale"),
#         gr.Number(value=1234, label="Seed", precision=0),
#     ],
#     outputs=[
#         gr.Image(label="Sample image"),
#         gr.Textbox(label="Sample text"),
#     ],
# )
# iface.launch()