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) |