from pydoc import describe
import gradio as gr
import torch
from omegaconf import OmegaConf
import sys
sys.path.append(".")
sys.path.append('./taming-transformers')
sys.path.append('./latent-diffusion')
from taming.models import vqgan
from ldm.util import instantiate_from_config
from huggingface_hub import hf_hub_download
import spaces
model_path_e = hf_hub_download(repo_id="multimodalart/compvis-latent-diffusion-text2img-large", filename="txt2img-f8-large.ckpt")
#@title Import stuff
import argparse, os, sys, glob
import numpy as np
from PIL import Image
from einops import rearrange
from torchvision.utils import make_grid
import transformers
import gc
from ldm.util import instantiate_from_config
from ldm.models.diffusion.ddim import DDIMSampler
from ldm.models.diffusion.plms import PLMSSampler
def load_model_from_config(config, ckpt, verbose=False):
print(f"Loading model from {ckpt}")
pl_sd = torch.load(ckpt, map_location="cpu")
sd = pl_sd["state_dict"]
model = instantiate_from_config(config.model)
m, u = model.load_state_dict(sd, strict=False)
if len(m) > 0 and verbose:
print("missing keys:")
print(m)
if len(u) > 0 and verbose:
print("unexpected keys:")
print(u)
del sd, pl_sd
gc.collect()
return model
config = OmegaConf.load("latent-diffusion/configs/latent-diffusion/txt2img-1p4B-eval.yaml")
model = load_model_from_config(config,model_path_e)
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
model = model.to(device)
@spaces.GPU(duration=120)
def run(prompt, steps, width, height, images, scale, progress=gr.Progress(track_tqdm=True)):
print(f"[run] entered with prompt={prompt!r} steps={steps} W={width} H={height} n={images} scale={scale}", flush=True)
opt = argparse.Namespace(
prompt = prompt,
outdir='latent-diffusion/outputs',
ddim_steps = int(steps),
ddim_eta = 0,
n_iter = 1,
W=int(width),
H=int(height),
n_samples=int(images),
scale=scale,
plms=True
)
print(f"[run] model device check: next(model.parameters()).device = {next(model.parameters()).device}", flush=True)
if opt.plms:
opt.ddim_eta = 0
sampler = PLMSSampler(model)
else:
sampler = DDIMSampler(model)
print("[run] sampler ready", flush=True)
os.makedirs(opt.outdir, exist_ok=True)
outpath = opt.outdir
prompt = opt.prompt
sample_path = os.path.join(outpath, "samples")
os.makedirs(sample_path, exist_ok=True)
base_count = len(os.listdir(sample_path))
all_samples=list()
all_samples_images=list()
with torch.no_grad():
with torch.cuda.amp.autocast():
with model.ema_scope():
uc = None
if opt.scale > 0:
print("[run] computing unconditional conditioning...", flush=True)
uc = model.get_learned_conditioning(opt.n_samples * [""])
print("[run] unconditional conditioning done", flush=True)
for n in range(opt.n_iter):
print(f"[run] computing prompt conditioning for iter {n}...", flush=True)
c = model.get_learned_conditioning(opt.n_samples * [prompt])
print("[run] prompt conditioning done; starting sampler.sample", flush=True)
shape = [4, opt.H//8, opt.W//8]
samples_ddim, _ = sampler.sample(S=opt.ddim_steps,
conditioning=c,
batch_size=opt.n_samples,
shape=shape,
verbose=False,
unconditional_guidance_scale=opt.scale,
unconditional_conditioning=uc,
eta=opt.ddim_eta)
x_samples_ddim = model.decode_first_stage(samples_ddim)
x_samples_ddim = torch.clamp((x_samples_ddim+1.0)/2.0, min=0.0, max=1.0)
for x_sample in x_samples_ddim:
x_sample = 255. * rearrange(x_sample.cpu().numpy(), 'c h w -> h w c')
image_vector = Image.fromarray(x_sample.astype(np.uint8))
all_samples_images.append(image_vector)
base_count += 1
all_samples.append(x_samples_ddim)
# additionally, save as grid
grid = torch.stack(all_samples, 0)
grid = rearrange(grid, 'n b c h w -> (n b) c h w')
grid = make_grid(grid, nrow=2)
# to image
grid = 255. * rearrange(grid, 'c h w -> h w c').cpu().numpy()
Image.fromarray(grid.astype(np.uint8)).save(os.path.join(outpath, f'{prompt.replace(" ", "-")}.png'))
return(all_samples_images,Image.fromarray(grid.astype(np.uint8)),None)
image = gr.Image(type="pil", label="Image Grid")
css = ".output-image{height: 528px !important} .output-carousel .output-image{height:272px !important} a{text-decoration: underline}"
iface = gr.Interface(fn=run, inputs=[
gr.Textbox(label="Prompt - try adding increments to your prompt such as 'oil on canvas', 'a painting', 'a book cover'",placeholder="chalk pastel drawing of a dog wearing a funny hat"),
gr.Slider(label="Steps - more steps can increase quality but will take longer to generate",value=45,maximum=50,minimum=1,step=1),
gr.Radio(label="Width", choices=[32,64,128,256],value=256),
gr.Radio(label="Height", choices=[32,64,128,256],value=256),
gr.Slider(label="Images - How many images you wish to generate", value=2, step=1, minimum=1, maximum=4),
gr.Slider(label="Diversity scale - How different from one another you wish the images to be",value=5.0, minimum=1.0, maximum=15.0),
#gr.inputs.Slider(label="ETA - between 0 and 1. Lower values can provide better quality, higher values can be more diverse",default=0.0,minimum=0.0, maximum=1.0,step=0.1),
],
outputs=[
gr.Gallery(label="Individual images"),
image,
gr.Textbox(label="Error")
],
css=css,
title="Generate images from text with Latent Diffusion LAION-400M",
description="",
article="Biases acknowledgment
Despite how impressive being able to turn text into image is, beware to the fact that this model may output content that reinforces or exarcbates societal biases. According to the
Latent Diffusion paper:
\"Deep learning modules tend to reproduce or exacerbate biases that are already present in the data\". The model was trained on an unfiltered version the LAION-400M dataset, which scrapped non-curated image-text-pairs from the internet (the exception being the the removal of illegal content) and is meant to be used for research purposes, such as this one.
You can read more on LAION's websiteWho owns the images produced by this demo?
")
iface.launch()