"""See https://huggingface.co/fu sing/latent-diffusion-text2im-large."""
from logzero import logger
from install import install
try:
    import gradio as gr
except ModuleNotFoundError:
        try:
            install("gradio")
            import gradio as gr
        except Exception as exc:
            logger.error(exc)
            raise SystemExit(1)

import PIL
from diffusers import DiffusionPipeline

ldm = DiffusionPipeline.from_pretrained("fu sing/latent-diffusion-text2im-large")

generator = torch.manual_seed(42)

examples = ["A street sign that reads Huggingface", "A painting of a squirrel eating a burger"]

prompt_ = "A painting of a squirrel eating a burger"

def fn(prompt=prompt_):
  image = ldm(
    [prompt],
    generator=generator,
    eta=0.3,
    guidance_scale=6.0,
    num_inference_steps=50,
  )

  image_processed = image.cpu().permute(0, 2, 3, 1)
  image_processed = image_processed  * 255.
  image_processed = image_processed.numpy().astype(np.uint8)
  image_pil = PIL.Image.fromarray(image_processed[0])

  # save image
  # image_pil.save("test.png")
  return image_pil

iface = gr.Interface(
  fn=fn,
  inputs="text",
  outputs="image",
  examples=examples,
  live=True,
)
iface.launch()

# gr.Interface.load("fu sing/latent-diffusion-text2im-large", examples=examples).launch()