stylemc.py

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