app.py

import gradio as gr
import torch
from diffusers import DiffusionPipeline

print(f"Is CUDA available: {torch.cuda.is_available()}")

if torch.cuda.is_available():
    print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
    pipe_vq = DiffusionPipeline.from_pretrained("microsoft/vq-diffusion-ithq", torch_dtype=torch.float16, revision="fp16").to("cuda")
else:
    pipe_vq = DiffusionPipeline.from_pretrained("microsoft/vq-diffusion-ithq")

examples = [
  ["An astronaut riding a horse."],
  ["A teddy bear playing in the water."],
  ["A simple wedding cake with lego bride and groom topper and cake pops."],
  ["A realistic tree using a mixture of different colored pencils."],
  ["Muscular Santa Claus."],
  ["A man with a pineapple head."],
  ["Pebble tower standing on the left on the sea beach."],
]

title = "VQ Diffusion vs. Stable Diffusion 1-5"
description = "[VQ-Diffusion-ITHQ](https://huggingface.co/microsoft/vq-diffusion-ithq) for text to image generation."


def inference(text):
    output_vq_diffusion = pipe_vq(text, truncation_rate=0.86).images[0]
    return output_vq_diffusion


io = gr.Interface(
  inference,
  gr.Textbox(lines=3),
  outputs=[
    gr.Image(type="pil", label="VQ-Diffusion"),
  ],
  title=title,
  description=description,
  examples=examples
)
io.launch()