Spaces:
Runtime error
Runtime error
from diffusers import AutoencoderKL, UNet2DConditionModel, DDIMScheduler | |
from transformers import CLIPTextModel, CLIPTokenizer, logging | |
import torch | |
from torchvision import transforms as tfms | |
from tqdm.auto import tqdm | |
from PIL import Image | |
# Supress some unnecessary warnings when loading the CLIPTextModel | |
logging.set_verbosity_error() | |
# Set device | |
device = "cuda" if torch.cuda.is_available() else "cpu" | |
# Loading components we'll use | |
tokenizer = CLIPTokenizer.from_pretrained( | |
"openai/clip-vit-large-patch14", | |
) | |
text_encoder = CLIPTextModel.from_pretrained( | |
"openai/clip-vit-large-patch14", | |
).to(device) | |
vae = AutoencoderKL.from_pretrained( | |
"CompVis/stable-diffusion-v1-4", | |
subfolder = "vae", | |
).to(device) | |
unet = UNet2DConditionModel.from_pretrained( | |
"CompVis/stable-diffusion-v1-4", | |
subfolder = "unet", | |
).to(device) | |
beta_start,beta_end = 0.00085,0.012 | |
scheduler = DDIMScheduler( | |
beta_start=beta_start, | |
beta_end=beta_end, | |
beta_schedule="scaled_linear", | |
num_train_timesteps=1000, | |
clip_sample=False, | |
set_alpha_to_one=False, | |
) | |
# convert PIL image to latents | |
def encode(img): | |
with torch.no_grad(): | |
latent = vae.encode(tfms.ToTensor()(img).unsqueeze(0).to(device)*2-1) | |
latent = 0.18215 * latent.latent_dist.sample() | |
return latent | |
# convert latents to PIL image | |
def decode(latent): | |
latent = (1 / 0.18215) * latent | |
with torch.no_grad(): | |
img = vae.decode(latent).sample | |
img = (img / 2 + 0.5).clamp(0, 1) | |
img = img.detach().cpu().permute(0, 2, 3, 1).numpy() | |
img = (img * 255).round().astype("uint8") | |
return Image.fromarray(img[0]) | |
# convert prompt into text embeddings, also unconditional embeddings | |
def prep_text(prompt): | |
text_input = tokenizer( | |
prompt, | |
padding="max_length", | |
max_length=tokenizer.model_max_length, | |
truncation=True, | |
return_tensors="pt", | |
) | |
text_embedding = text_encoder( | |
text_input.input_ids.to(device) | |
)[0] | |
uncond_input = tokenizer( | |
"", | |
padding="max_length", | |
max_length=tokenizer.model_max_length, | |
truncation=True, | |
return_tensors="pt", | |
) | |
uncond_embedding = text_encoder( | |
uncond_input.input_ids.to(device) | |
)[0] | |
return torch.cat([uncond_embedding, text_embedding]) | |
def magic_mix( | |
img, # specifies the layout semantics | |
prompt, # specifies the content semantics | |
kmin=0.3, | |
kmax=0.6, | |
v=0.5, # interpolation constant | |
seed=42, | |
steps=50, | |
guidance_scale=7.5, | |
): | |
tmin = steps- int(kmin*steps) | |
tmax = steps- int(kmax*steps) | |
text_embeddings = prep_text(prompt) | |
scheduler.set_timesteps(steps) | |
width, height = img.size | |
encoded = encode(img) | |
torch.manual_seed(seed) | |
noise = torch.randn( | |
(1,unet.in_channels,height // 8,width // 8), | |
).to(device) | |
latents = scheduler.add_noise( | |
encoded, | |
noise, | |
timesteps=scheduler.timesteps[tmax] | |
) | |
input = torch.cat([latents]*2) | |
input = scheduler.scale_model_input(input, scheduler.timesteps[tmax]) | |
with torch.no_grad(): | |
pred = unet( | |
input, | |
scheduler.timesteps[tmax], | |
encoder_hidden_states=text_embeddings, | |
).sample | |
pred_uncond, pred_text = pred.chunk(2) | |
pred = pred_uncond + guidance_scale * (pred_text - pred_uncond) | |
latents = scheduler.step(pred, scheduler.timesteps[tmax], latents).prev_sample | |
for i, t in enumerate(tqdm(scheduler.timesteps)): | |
if i > tmax: | |
if i < tmin: # layout generation phase | |
orig_latents = scheduler.add_noise( | |
encoded, | |
noise, | |
timesteps=t | |
) | |
input = (v*latents) + (1-v)*orig_latents # interpolating between layout noise and conditionally generated noise to preserve layout sematics | |
input = torch.cat([input]*2) | |
else: # content generation phase | |
input = torch.cat([latents]*2) | |
input = scheduler.scale_model_input(input, t) | |
with torch.no_grad(): | |
pred = unet( | |
input, | |
t, | |
encoder_hidden_states=text_embeddings, | |
).sample | |
pred_uncond, pred_text = pred.chunk(2) | |
pred = pred_uncond + guidance_scale * (pred_text - pred_uncond) | |
latents = scheduler.step(pred, t, latents).prev_sample | |
return decode(latents) | |
if __name__ == "__main__": | |
import argparse | |
parser = argparse.ArgumentParser() | |
parser.add_argument("img_file", type=str, help="image file to provide the layout semantics for the mixing process") | |
parser.add_argument("prompt", type=str, help="prompt to provide the content semantics for the mixing process") | |
parser.add_argument("out_file", type=str, help="filename to save the generation to") | |
parser.add_argument("--kmin", type=float, default=0.3) | |
parser.add_argument("--kmax", type=float, default=0.6) | |
parser.add_argument("--v", type=float, default=0.5) | |
parser.add_argument("--seed", type=int, default=42) | |
parser.add_argument("--steps", type=int, default=50) | |
parser.add_argument("--guidance_scale", type=float, default=7.5) | |
args = parser.parse_args() | |
img = Image.open(args.img_file) | |
out_img = magic_mix( | |
img, | |
args.prompt, | |
args.kmin, | |
args.kmax, | |
args.v, | |
args.seed, | |
args.steps, | |
args.guidance_scale | |
) | |
out_img.save(args.out_file) |