""" Approach: "StyleMC: Multi-Channel Based Fast Text-Guided Image Generation and Manipulation" Original source code: https://github.com/autonomousvision/stylegan_xl/blob/f9be58e98110bd946fcdadef2aac8345466faaf3/run_stylemc.py# Modified by Håkon Hukkelås """ import os from pathlib import Path import tqdm import re import click from dp2 import utils import tops from typing import List, Optional import PIL.Image import imageio from timeit import default_timer as timer import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from torchvision.transforms.functional import resize, normalize from dp2.infer import build_trained_generator import clip #---------------------------------------------------------------------------- class AverageMeter(object): """Computes and stores the average and current value""" def __init__(self, name, fmt=':f'): self.name = name self.fmt = fmt self.reset() def reset(self): self.val = 0 self.avg = 0 self.sum = 0 self.count = 0 def update(self, val, n=1): self.val = val self.sum += val * n self.count += n self.avg = self.sum / self.count def __str__(self): fmtstr = '{name} {val' + self.fmt + '} ({avg' + self.fmt + '})' return fmtstr.format(**self.__dict__) class ProgressMeter(object): def __init__(self, num_batches, meters, prefix=""): self.batch_fmtstr = self._get_batch_fmtstr(num_batches) self.meters = meters self.prefix = prefix def display(self, batch): entries = [self.prefix + self.batch_fmtstr.format(batch)] entries += [str(meter) for meter in self.meters] print('\t'.join(entries)) def _get_batch_fmtstr(self, num_batches): num_digits = len(str(num_batches // 1)) fmt = '{:' + str(num_digits) + 'd}' return '[' + fmt + '/' + fmt.format(num_batches) + ']' def save_image(img, path): img = (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8) PIL.Image.fromarray(img[0].cpu().numpy(), 'RGB').save(path) def unravel_index(index, shape): out = [] for dim in reversed(shape): out.append(index % dim) index = index // dim return tuple(reversed(out)) def num_range(s: str) -> List[int]: '''Accept either a comma separated list of numbers 'a,b,c' or a range 'a-c' and return as a list of ints.''' range_re = re.compile(r'^(\d+)-(\d+)$') m = range_re.match(s) if m: return list(range(int(m.group(1)), int(m.group(2))+1)) vals = s.split(',') return [int(x) for x in vals] #---------------------------------------------------------------------------- def spherical_dist_loss(x, y): x = F.normalize(x, dim=-1) y = F.normalize(y, dim=-1) return (x - y).norm(dim=-1).div(2).arcsin().pow(2).mul(2) def prompts_dist_loss(x, targets, loss): if len(targets) == 1: # Keeps consistent results vs previous method for single objective guidance return loss(x, targets[0]) distances = [loss(x, target) for target in targets] return torch.stack(distances, dim=-1).sum(dim=-1) def embed_text(model, prompt, device='cuda'): return #---------------------------------------------------------------------------- @torch.no_grad() @torch.cuda.amp.autocast() def generate_edit( G, dl, direction, edit_strength, path, ): for it, batch in enumerate(dl): batch["embedding"] = None styles = get_styles(None, G, batch, truncation_value=0) imgs = [] grad_changes = [_*edit_strength for _ in [0, 0.25, 0.5, 0.75, 1]] grad_changes = [*[-x for x in grad_changes][::-1], *grad_changes] batch = {k: tops.to_cuda(v) if v is not None else v for k,v in batch.items()} for i, grad_change in enumerate(grad_changes): s = styles + direction*grad_change img = G(**batch, s=iter(s))["img"] img = (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255) imgs.append(img[0].to(torch.uint8).cpu().numpy()) PIL.Image.fromarray(np.concatenate(imgs, axis=1), 'RGB').save(path + f'{it}.png') @torch.no_grad() def get_styles(seed, G: torch.nn.Module, batch, truncation_value=1): all_styles = [] if seed is None: z = np.random.normal(0, 0, size=(1, G.z_channels)) else: z = np.random.RandomState(seed=seed).normal(0, 1, size=(1, G.z_channels)) z_idx = np.random.RandomState(seed=seed).randint(0, len(G.style_net.w_centers)) w_c = G.style_net.w_centers[z_idx].to(tops.get_device()).view(1, -1) w = G.style_net(torch.from_numpy(z).to(tops.get_device())) w = w_c.to(w.dtype).lerp(w, truncation_value) if hasattr(G, "get_comod_y"): w = G.get_comod_y(batch, w) for block in G.modules(): if not hasattr(block, "affine") or not hasattr(block.affine, "weight"): continue gamma0 = block.affine(w) if hasattr(block, "affine_beta"): beta0 = block.affine_beta(w) gamma0 = torch.cat((gamma0, beta0), dim=1) all_styles.append(gamma0) max_ch = max([s.shape[-1] for s in all_styles]) all_styles = [F.pad(s, ((0, max_ch - s.shape[-1])), "constant", 0) for s in all_styles] all_styles = torch.cat(all_styles) return all_styles def get_and_cache_direction(output_dir: Path, dl_val, G, text_prompt): cache_path = output_dir.joinpath( "stylemc_cache", text_prompt.replace(" ", "_") + ".torch") if cache_path.is_file(): print("Loaded cache from:", cache_path) return torch.load(cache_path) direction = find_direction(G, text_prompt, None, dl_val=iter(dl_val)) cache_path.parent.mkdir(exist_ok=True, parents=True) torch.save(direction, cache_path) return direction @torch.cuda.amp.autocast() def find_direction( G, text_prompt, batches, #layers, n_iterations=128*8, batch_size=8, dl_val=None ): time_start = timer() clip_model = clip.load("ViT-B/16", device=tops.get_device())[0] target = [clip_model.encode_text(clip.tokenize(text_prompt).to(tops.get_device())).float()] all_styles = [] if dl_val is not None: first_batch = next(dl_val) else: first_batch = batches[0] first_batch["embedding"] = None if "embedding" not in first_batch else first_batch["embedding"] s = get_styles(0, G, first_batch) # stats tracker cos_sim_track = AverageMeter('cos_sim', ':.4f') norm_track = AverageMeter('norm', ':.4f') n_iterations = n_iterations // batch_size progress = ProgressMeter(n_iterations, [cos_sim_track, norm_track]) # initalize styles direction direction = torch.zeros(s.shape, device=tops.get_device()) direction.requires_grad_() utils.set_requires_grad(G, False) direction_tracker = torch.zeros_like(direction) opt = torch.optim.AdamW([direction], lr=0.05, betas=(0., 0.999), weight_decay=0.25) grads = [] for seed_idx in tqdm.trange(n_iterations): # forward pass through synthesis network with new styles if seed_idx == 0: batch = first_batch elif dl_val is not None: batch = next(dl_val) batch["embedding"] = None if "embedding" not in batch else batch["embedding"] else: batch = {k: tops.to_cuda(v) if v is not None else v for k, v in batches[seed_idx].items()} styles = get_styles(seed_idx, G, batch) + direction img = G(**batch, s=iter(styles))["img"] batch = {k: v.cpu() if v is not None else v for k, v in batch.items()} # clip loss img = (img + 1)/2 img = normalize(img, mean=(0.48145466, 0.4578275, 0.40821073), std=(0.26862954, 0.26130258, 0.27577711)) img = resize(img, (224, 224)) embeds = clip_model.encode_image(img) cos_sim = prompts_dist_loss(embeds, target, spherical_dist_loss) cos_sim.backward(retain_graph=True) # track stats cos_sim_track.update(cos_sim.item()) norm_track.update(torch.norm(direction).item()) if not (seed_idx % batch_size): # zeroing out gradients for non-optimized layers #layers_zeroed = torch.tensor([x for x in range(G.num_ws) if not x in layers]) #direction.grad[:, layers_zeroed] = 0 opt.step() grads.append(direction.grad.clone()) direction.grad.data.zero_() # keep track of gradients over time if seed_idx > 3: direction_tracker[grads[-2] * grads[-1] < 0] += 1 # plot stats progress.display(seed_idx) # throw out fluctuating channels direction = direction.detach() direction[direction_tracker > n_iterations / 4] = 0 print(direction) print(f"Time for direction search: {timer() - time_start:.2f} s") return direction @click.command() @click.argument("config_path") @click.argument("input_path") @click.argument("output_path") #@click.option('--layers', type=num_range, help='Restrict the style space to a range of layers. We recommend not to optimize the critically sampled layers (last 3).', required=True) @click.option('--text-prompt', help='Text', type=str, required=True) @click.option('--edit-strength', help='Strength of edit', type=float, required=True) @click.option('--outdir', help='Where to save the output images', type=str, required=True) def stylemc( config_path, #layers: List[int], text_prompt: str, edit_strength: float, outdir: str, ): cfg = utils.load_config(config_path) G = build_trained_generator(cfg) cfg.train.batch_size = 1 n_iterations = 256 dl_val = tops.config.instantiate(cfg.data.val.loader) direction = find_direction(G, text_prompt, None, n_iterations=n_iterations, dl_val=iter(dl_val)) text_prompt = text_prompt.replace(" ", "_") generate_edit(G, input_path, direction, edit_strength, output_path) if __name__ == "__main__": stylemc()