update demo

generate_fromS.py

@@ -1,277 +0,0 @@
-# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
-#
-# NVIDIA CORPORATION and its licensors retain all intellectual property
-# and proprietary rights in and to this software, related documentation
-# and any modifications thereto.  Any use, reproduction, disclosure or
-# distribution of this software and related documentation without an express
-# license agreement from NVIDIA CORPORATION is strictly prohibited.
-
-"""Generate images using pretrained network pickle."""
-
-import os
-import re
-import random
-import math
-import time
-import click
-import legacy
-from typing import List, Optional
-
-import cv2
-import clip
-import dnnlib
-import numpy as np
-import torchvision
-from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize
-import PIL.Image
-import matplotlib.pyplot as plt
-import torch
-from torch import linalg as LA
-import torch.nn.functional as F
-from torch_utils import misc
-from torch_utils import persistence
-from torch_utils.ops import conv2d_resample
-from torch_utils.ops import upfirdn2d
-from torch_utils.ops import bias_act
-from torch_utils.ops import fma
-
-
-def block_forward(self, x, img, ws, shapes, force_fp32=False, fused_modconv=None, **layer_kwargs):
-        misc.assert_shape(ws, [None, self.num_conv + self.num_torgb, self.w_dim])
-        w_iter = iter(ws.unbind(dim=1))
-        dtype = torch.float16 if self.use_fp16 and not force_fp32 else torch.float32
-        memory_format = torch.channels_last if self.channels_last and not force_fp32 else torch.contiguous_format
-        if fused_modconv is None:
-            with misc.suppress_tracer_warnings(): # this value will be treated as a constant
-                fused_modconv = (not self.training) and (dtype == torch.float32 or int(x.shape[0]) == 1)
-
-        # Input.
-        if self.in_channels == 0:
-            x = self.const.to(dtype=dtype, memory_format=memory_format)
-            x = x.unsqueeze(0).repeat([ws.shape[0], 1, 1, 1])
-        else:
-            misc.assert_shape(x, [None, self.in_channels, self.resolution // 2, self.resolution // 2])
-            x = x.to(dtype=dtype, memory_format=memory_format)
-
-        # Main layers.
-        if self.in_channels == 0:
-            x = self.conv1(x, next(w_iter)[...,:shapes[0]], fused_modconv=fused_modconv, **layer_kwargs)
-        elif self.architecture == 'resnet':
-            y = self.skip(x, gain=np.sqrt(0.5))
-            x = self.conv0(x, next(w_iter), fused_modconv=fused_modconv, **layer_kwargs)
-            x = self.conv1(x, next(w_iter), fused_modconv=fused_modconv, gain=np.sqrt(0.5), **layer_kwargs)
-            x = y.add_(x)
-        else:
-            x = self.conv0(x, next(w_iter)[...,:shapes[0]], fused_modconv=fused_modconv, **layer_kwargs)
-            x = self.conv1(x, next(w_iter)[...,:shapes[1]], fused_modconv=fused_modconv, **layer_kwargs)
-
-        # ToRGB.
-        if img is not None:
-            misc.assert_shape(img, [None, self.img_channels, self.resolution // 2, self.resolution // 2])
-            img = upfirdn2d.upsample2d(img, self.resample_filter)
-        if self.is_last or self.architecture == 'skip':
-            y = self.torgb(x, next(w_iter)[...,:shapes[2]], fused_modconv=fused_modconv)
-            y = y.to(dtype=torch.float32, memory_format=torch.contiguous_format)
-            img = img.add_(y) if img is not None else y
-
-        assert x.dtype == dtype
-        assert img is None or img.dtype == torch.float32
-        return x, img
-
-
-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]
-
-
-@click.command()
-@click.pass_context
-@click.option('--network', 'network_pkl', help='Network pickle filename', required=True)
-@click.option('--seeds', type=num_range, help='List of random seeds')
-@click.option('--trunc', 'truncation_psi', type=float, help='Truncation psi', default=0.7, show_default=True)
-@click.option('--class', 'class_idx', type=int, help='Class label (unconditional if not specified)')
-@click.option('--noise-mode', help='Noise mode', type=click.Choice(['const', 'random', 'none']), default='const', show_default=True)
-@click.option('--projected-w', help='Projection result file', type=str, metavar='FILE')
-@click.option('--s_input', help='Projection result file', type=str, metavar='FILE')
-@click.option('--outdir', help='Where to save the output images', type=str, required=True, metavar='DIR')
-@click.option('--text_prompt', help='Text', type=str, required=True)
-@click.option('--change_power', help='Change power', type=int, required=True)
-@click.option('--from_video', 'from_video', is_flag=True, help="generate from video")
-
-def generate_images(
-    ctx: click.Context,
-    network_pkl: str,
-    seeds: Optional[List[int]],
-    truncation_psi: float,
-    noise_mode: str,
-    outdir: str,
-    class_idx: Optional[int],
-    projected_w: Optional[str],
-    s_input: Optional[str],
-    text_prompt: str,
-    change_power: int,
-    from_video: bool,
-):
-    """
-    Generate images using pretrained network pickle.
-
-    Examples:
-    # Generate curated MetFaces images without truncation (Fig.10 left)
-    python generate.py --outdir=out --trunc=1 --seeds=85,265,297,849 \\
-        --network=https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/metfaces.pkl
-
-    # Generate uncurated MetFaces images with truncation (Fig.12 upper left)
-    python generate.py --outdir=out --trunc=0.7 --seeds=600-605 \\
-        --network=https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/metfaces.pkl
-
-    # Generate class conditional CIFAR-10 images (Fig.17 left, Car)
-    python generate.py --outdir=out --seeds=0-35 --class=1 \\
-        --network=https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/cifar10.pkl
-
-    # Render an image from projected W
-    python generate.py --outdir=out --projected_w=projected_w.npz \\
-        --network=https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/metfaces.pkl
-    """
-
-    print('Loading networks from "%s"...' % network_pkl)
-    # Use GPU if available
-    if torch.cuda.is_available():
-        device = torch.device("cuda")
-    else:
-        device = torch.device("cpu")
-
-    with dnnlib.util.open_url(network_pkl) as f:
-        G = legacy.load_network_pkl(f)['G_ema'].to(device) # type: ignore
-
-    os.makedirs(outdir, exist_ok=True)
-
-    # Synthesize the result of a W projection.
-    if projected_w is not None:
-        if seeds is not None:
-            print ('warn: --seeds is ignored when using --projected-w')
-        print(f'Generating images from projected W "{projected_w}"')
-        ws = np.load(projected_w)['w']
-        ws = torch.tensor(ws, device=device) # pylint: disable=not-callable
-        assert ws.shape[1:] == (G.num_ws, G.w_dim)
-        for idx, w in enumerate(ws):
-            img = G.synthesis(w.unsqueeze(0), noise_mode=noise_mode)
-            img = (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8)
-            img = PIL.Image.fromarray(img[0].cpu().numpy(), 'RGB')
-            img.save(f'{outdir}/proj{idx:02d}.png')
-        return
-
-    # Labels
-    label = torch.zeros([1, G.c_dim], device=device).requires_grad_()
-    if G.c_dim != 0:
-        if class_idx is None:
-            ctx.fail('Must specify class label with --class when using a conditional network')
-        label[:, class_idx] = 1
-    else:
-        if class_idx is not None:
-            print ('warn: --class=lbl ignored when running on an unconditional network')
-
-    # Generate images
-    for i in G.parameters():
-      i.requires_grad = False
-
-
-    temp_shapes = []
-    for res in G.synthesis.block_resolutions:
-        block = getattr(G.synthesis, f'b{res}')
-        if res == 4:
-            temp_shape = (block.conv1.affine.weight.shape[0], block.conv1.affine.weight.shape[0], block.torgb.affine.weight.shape[0])
-            block.conv1.affine = torch.nn.Identity()
-            block.torgb.affine = torch.nn.Identity()
-
-        else:
-            temp_shape = (block.conv0.affine.weight.shape[0], block.conv1.affine.weight.shape[0], block.torgb.affine.weight.shape[0])
-            block.conv0.affine = torch.nn.Identity()
-            block.conv1.affine = torch.nn.Identity()
-            block.torgb.affine = torch.nn.Identity()
-
-        temp_shapes.append(temp_shape)
-
-
-    if s_input is not None:
-        styles = np.load(s_input)['s']
-        styles_direction = np.load(f'{outdir}/direction_'+text_prompt.replace(" ", "_")+'.npz')['s']
-
-        styles_direction = torch.tensor(styles_direction, device=device)
-        styles = torch.tensor(styles, device=device)
-
-    if from_video and not os.path.isdir(f'{outdir}_video'):
-        os.makedirs(f'{outdir}_video')
-
-    with torch.no_grad():
-        if from_video:
-            name_i = 1000
-            for grad_change in np.arange(0, 1, 0.02)*change_power:
-                imgs = []
-                name_i += 1
-
-                styles += styles_direction*grad_change
-                styles_idx = 0
-                x = img = None
-                for k , res in enumerate(G.synthesis.block_resolutions):
-                    block = getattr(G.synthesis, f'b{res}')
-
-                    if res == 4:
-                        x, img = block_forward(block, x, img, styles[:, styles_idx:styles_idx+2, :], temp_shapes[k], noise_mode=noise_mode, force_fp32=True)
-                        styles_idx += 2
-                    else:
-                        x, img = block_forward(block, x, img, styles[:, styles_idx:styles_idx+3, :], temp_shapes[k], noise_mode=noise_mode, force_fp32=True)
-                        styles_idx += 3
-
-                img = (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255)
-                imgs.append(img[0].to(torch.uint8).cpu().numpy())
-
-                styles -= styles_direction*grad_change
-                img_filepath = '{}_video/{}_{}_{}.jpeg'.format(outdir, text_prompt.replace(" ", "_"), change_power, name_i)
-                PIL.Image.fromarray(np.concatenate(imgs, axis=1), 'RGB').save(img_filepath, quality=95)
-        else:
-            imgs = []
-            grad_changes = [0, 0.25*change_power, 0.5*change_power, 0.75*change_power, change_power]
-
-            for grad_change in grad_changes:
-                styles += styles_direction*grad_change
-
-                styles_idx = 0
-                x = img = None
-                for k , res in enumerate(G.synthesis.block_resolutions):
-                    block = getattr(G.synthesis, f'b{res}')
-
-                    if res == 4:
-                        x, img = block_forward(block, x, img, styles[:, styles_idx:styles_idx+2, :], temp_shapes[k], noise_mode=noise_mode, force_fp32=True)
-                        styles_idx += 2
-                    else:
-                        x, img = block_forward(block, x, img, styles[:, styles_idx:styles_idx+3, :], temp_shapes[k], noise_mode=noise_mode, force_fp32=True)
-                        styles_idx += 3
-
-                img = (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255)
-                imgs.append(img[0].to(torch.uint8).cpu().numpy())
-
-                styles -= styles_direction*grad_change
-
-            img_filepath = f'{outdir}/'+text_prompt.replace(" ", "_")+'_'+str(change_power)+'.jpeg'
-            PIL.Image.fromarray(np.concatenate(imgs, axis=1), 'RGB').save(img_filepath, quality=95)
-
-
-
-if __name__ == "__main__":
-    generate_images()

generate_multi.py

@@ -1,403 +0,0 @@
-# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
-#
-# NVIDIA CORPORATION and its licensors retain all intellectual property
-# and proprietary rights in and to this software, related documentation
-# and any modifications thereto.  Any use, reproduction, disclosure or
-# distribution of this software and related documentation without an express
-# license agreement from NVIDIA CORPORATION is strictly prohibited.
-
-"""Generate images using pretrained network pickle."""
-
-import os
-import re
-from typing import List, Optional
-import torchvision
-from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize
-import click
-import dnnlib
-import numpy as np
-import PIL.Image
-import torch
-from torch import linalg as LA
-import clip
-from PIL import Image
-import legacy
-import torch.nn.functional as F
-import cv2
-import matplotlib.pyplot as plt
-from torch_utils import misc
-from torch_utils import persistence
-from torch_utils.ops import conv2d_resample
-from torch_utils.ops import upfirdn2d
-from torch_utils.ops import bias_act
-from torch_utils.ops import fma
-import random
-import math
-import time
-import id_loss
-
-
-def block_forward(self, x, img, ws, shapes, force_fp32=False, fused_modconv=None, **layer_kwargs):
-        misc.assert_shape(ws, [None, self.num_conv + self.num_torgb, self.w_dim])
-        w_iter = iter(ws.unbind(dim=1))
-        dtype = torch.float16 if self.use_fp16 and not force_fp32 else torch.float32
-        memory_format = torch.channels_last if self.channels_last and not force_fp32 else torch.contiguous_format
-        if fused_modconv is None:
-            with misc.suppress_tracer_warnings(): # this value will be treated as a constant
-                fused_modconv = (not self.training) and (dtype == torch.float32 or int(x.shape[0]) == 1)
-
-        # Input.
-        if self.in_channels == 0:
-            x = self.const.to(dtype=dtype, memory_format=memory_format)
-            x = x.unsqueeze(0).repeat([ws.shape[0], 1, 1, 1])
-        else:
-            misc.assert_shape(x, [None, self.in_channels, self.resolution // 2, self.resolution // 2])
-            x = x.to(dtype=dtype, memory_format=memory_format)
-
-        # Main layers.
-        if self.in_channels == 0:
-            x = self.conv1(x, next(w_iter)[...,:shapes[0]], fused_modconv=fused_modconv, **layer_kwargs)
-        elif self.architecture == 'resnet':
-            y = self.skip(x, gain=np.sqrt(0.5))
-            x = self.conv0(x, next(w_iter), fused_modconv=fused_modconv, **layer_kwargs)
-            x = self.conv1(x, next(w_iter), fused_modconv=fused_modconv, gain=np.sqrt(0.5), **layer_kwargs)
-            x = y.add_(x)
-        else:
-            x = self.conv0(x, next(w_iter)[...,:shapes[0]], fused_modconv=fused_modconv, **layer_kwargs)
-            x = self.conv1(x, next(w_iter)[...,:shapes[1]], fused_modconv=fused_modconv, **layer_kwargs)
-
-        # ToRGB.
-        if img is not None:
-            misc.assert_shape(img, [None, self.img_channels, self.resolution // 2, self.resolution // 2])
-            img = upfirdn2d.upsample2d(img, self.resample_filter)
-        if self.is_last or self.architecture == 'skip':
-            y = self.torgb(x, next(w_iter)[...,:shapes[2]], fused_modconv=fused_modconv)
-            y = y.to(dtype=torch.float32, memory_format=torch.contiguous_format)
-            img = img.add_(y) if img is not None else y
-
-        assert x.dtype == dtype
-        assert img is None or img.dtype == torch.float32
-        return x, img
-
-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]
-
-#----------------------------------------------------------------------------
-
-@click.command()
-@click.pass_context
-@click.option('--network', 'network_pkl', help='Network pickle filename', required=True)
-@click.option('--seeds', type=num_range, help='List of random seeds')
-@click.option('--trunc', 'truncation_psi', type=float, help='Truncation psi', default=1, show_default=True)
-@click.option('--class', 'class_idx', type=int, help='Class label (unconditional if not specified)')
-@click.option('--noise-mode', help='Noise mode', type=click.Choice(['const', 'random', 'none']), default='const', show_default=True)
-@click.option('--projected-w', help='Projection result file', type=str, metavar='FILE')
-@click.option('--projected_s', help='Projection result file', type=str, metavar='FILE')
-@click.option('--outdir', help='Where to save the output images', type=str, required=True, metavar='DIR')
-@click.option('--resolution', help='Resolution of output images', type=int, required=True)
-@click.option('--batch_size', help='Batch Size', type=int, required=True)
-@click.option('--identity_power', help='How much change occurs on the face', type=str, required=True)
-def generate_images(
-    ctx: click.Context,
-    network_pkl: str,
-    seeds: Optional[List[int]],
-    truncation_psi: float,
-    noise_mode: str,
-    outdir: str,
-    class_idx: Optional[int],
-    projected_w: Optional[str],
-    projected_s: Optional[str],
-    resolution: int,
-    batch_size: int,
-    identity_power: str
-):
-    """Generate images using pretrained network pickle.
-
-    Examples:
-
-    \b
-    # Generate curated MetFaces images without truncation (Fig.10 left)
-    python generate.py --outdir=out --trunc=1 --seeds=85,265,297,849 \\
-        --network=https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/metfaces.pkl
-
-    \b
-    # Generate uncurated MetFaces images with truncation (Fig.12 upper left)
-    python generate.py --outdir=out --trunc=0.7 --seeds=600-605 \\
-        --network=https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/metfaces.pkl
-
-    \b
-    # Generate class conditional CIFAR-10 images (Fig.17 left, Car)
-    python generate.py --outdir=out --seeds=0-35 --class=1 \\
-        --network=https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/cifar10.pkl
-
-    \b
-    # Render an image from projected W
-    python generate.py --outdir=out --projected_w=projected_w.npz \\
-        --network=https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/metfaces.pkl
-    """
-
-    print('Loading networks from "%s"...' % network_pkl)
-    device = torch.device('cuda')
-    with dnnlib.util.open_url(network_pkl) as f:
-        G = legacy.load_network_pkl(f)['G_ema'].to(device) # type: ignore
-
-    os.makedirs(outdir, exist_ok=True)
-
-    # Synthesize the result of a W projection.
-    if projected_w is not None:
-        if seeds is not None:
-            print ('warn: --seeds is ignored when using --projected-w')
-        print(f'Generating images from projected W "{projected_w}"')
-        ws = np.load(projected_w)['w']
-        ws = torch.tensor(ws, device=device) # pylint: disable=not-callable
-        assert ws.shape[1:] == (G.num_ws, G.w_dim)
-        for idx, w in enumerate(ws):
-            img = G.synthesis(w.unsqueeze(0), noise_mode=noise_mode)
-            img = (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8)
-            img = PIL.Image.fromarray(img[0].cpu().numpy(), 'RGB').save(f'{outdir}/proj{idx:02d}.png')
-        return
-
-    if seeds is None:
-        ctx.fail('--seeds option is required when not using --projected-w')
-
-    # Labels.
-    label = torch.zeros([1, G.c_dim], device=device).requires_grad_()
-    if G.c_dim != 0:
-        if class_idx is None:
-            ctx.fail('Must specify class label with --class when using a conditional network')
-        label[:, class_idx] = 1
-    else:
-        if class_idx is not None:
-            print ('warn: --class=lbl ignored when running on an unconditional network')
-
-    model, preprocess = clip.load("ViT-B/32", device=device)
-
-    text_prompts_file = open("text_prompts.txt")
-    text_prompts = text_prompts_file.read().split("\n")
-    text_prompts_file.close()
-
-    text = clip.tokenize(text_prompts).to(device)
-    text_features = model.encode_text(text)
-
-    # Generate images.
-    for i in G.parameters():
-      i.requires_grad = True
-
-    mean = torch.as_tensor((0.48145466, 0.4578275, 0.40821073), dtype=torch.float, device=device)
-    std = torch.as_tensor((0.26862954, 0.26130258, 0.27577711), dtype=torch.float, device=device)
-    if mean.ndim == 1:
-        mean = mean.view(-1, 1, 1)
-    if std.ndim == 1:
-        std = std.view(-1, 1, 1)
-
-    transf = Compose([
-        Resize(224, interpolation=Image.BICUBIC),
-        CenterCrop(224),
-    ])
-
-    styles_array = []
-    print("seeds:", seeds)
-    t1 = time.time()
-    for seed_idx, seed in enumerate(seeds):
-        if seed==seeds[-1]:
-            print('Generating image for seed %d (%d/%d) ...' % (seed, seed_idx, len(seeds)))
-        z = torch.from_numpy(np.random.RandomState(seed).randn(1, G.z_dim)).to(device)
-        ws = G.mapping(z, label, truncation_psi=truncation_psi)
-
-        block_ws = []
-        with torch.autograd.profiler.record_function('split_ws'):
-            misc.assert_shape(ws, [None, G.synthesis.num_ws, G.synthesis.w_dim])
-            ws = ws.to(torch.float32)
-
-
-            w_idx = 0
-            for res in G.synthesis.block_resolutions:
-                block = getattr(G.synthesis, f'b{res}')
-                block_ws.append(ws.narrow(1, w_idx, block.num_conv + block.num_torgb))
-                w_idx += block.num_conv
-
-
-        styles = torch.zeros(1,26,512, device=device)
-        styles_idx = 0
-        temp_shapes = []
-        for res, cur_ws in zip(G.synthesis.block_resolutions, block_ws):
-            block = getattr(G.synthesis, f'b{res}')
-
-            if res == 4:
-              temp_shape = (block.conv1.affine.weight.shape[0], block.conv1.affine.weight.shape[0], block.torgb.affine.weight.shape[0])
-              styles[0,:1,:] = block.conv1.affine(cur_ws[0,:1,:])
-              styles[0,1:2,:] = block.torgb.affine(cur_ws[0,1:2,:])
-              if seed_idx==(len(seeds)-1):
-                block.conv1.affine = torch.nn.Identity()
-                block.torgb.affine = torch.nn.Identity()
-              styles_idx += 2
-            else:
-              temp_shape = (block.conv0.affine.weight.shape[0], block.conv1.affine.weight.shape[0], block.torgb.affine.weight.shape[0])
-              styles[0,styles_idx:styles_idx+1,:temp_shape[0]] = block.conv0.affine(cur_ws[0,:1,:])
-              styles[0,styles_idx+1:styles_idx+2,:temp_shape[1]] = block.conv1.affine(cur_ws[0,1:2,:])
-              styles[0,styles_idx+2:styles_idx+3,:temp_shape[2]] = block.torgb.affine(cur_ws[0,2:3,:])
-              if seed_idx==(len(seeds)-1):
-                block.conv0.affine = torch.nn.Identity()
-                block.conv1.affine = torch.nn.Identity()
-                block.torgb.affine = torch.nn.Identity()
-              styles_idx += 3
-            temp_shapes.append(temp_shape)
-
-
-        styles = styles.detach()
-        styles_array.append(styles)
-
-    resolution_dict = {256: 6, 512: 7, 1024: 8}
-    identity_coefficient_dict = {"high": 2,"medium": 0.5, "low": 0.1, "none": 0}
-    identity_coefficient = identity_coefficient_dict[identity_power]
-    styles_wanted_direction = torch.zeros(1,26,512, device=device)
-    styles_wanted_direction_grad_el2 = torch.zeros(1,26,512, device=device)
-    styles_wanted_direction.requires_grad_()
-
-    global id_loss
-    id_loss = id_loss.IDLoss("a").to(device).eval()
-
-    temp_photos = []
-    grads = []
-    for i in range(math.ceil(len(seeds)/batch_size)):
-        #print(i*batch_size, "processed", time.time()-t1)
-
-
-        styles = torch.vstack(styles_array[i*batch_size:(i+1)*batch_size]).to(device)
-
-
-        seed = seeds[i]
-
-        styles_idx = 0
-        x2 = img2 = None
-
-        for k , (res, cur_ws) in enumerate(zip(G.synthesis.block_resolutions, block_ws)):
-            block = getattr(G.synthesis, f'b{res}')
-            if k>resolution_dict[resolution]:
-              continue
-
-            if res == 4:
-              x2, img2 = block_forward(block, x2, img2, styles[:, styles_idx:styles_idx+2, :], temp_shapes[k], noise_mode=noise_mode)
-              styles_idx += 2
-            else:
-              x2, img2 = block_forward(block, x2, img2, styles[:, styles_idx:styles_idx+3, :], temp_shapes[k], noise_mode=noise_mode)
-              styles_idx += 3
-
-        img2_cpu = img2.detach().cpu().numpy()
-        temp_photos.append(img2_cpu)
-        if i>3:
-            continue
-
-        styles2 = styles + styles_wanted_direction
-
-        styles_idx = 0
-        x = img = None
-        for k , (res, cur_ws) in enumerate(zip(G.synthesis.block_resolutions, block_ws)):
-            block = getattr(G.synthesis, f'b{res}')
-            if k>resolution_dict[resolution]:
-              continue
-            if res == 4:
-              x, img = block_forward(block, x, img, styles2[:, styles_idx:styles_idx+2, :], temp_shapes[k], noise_mode=noise_mode)
-              styles_idx += 2
-            else:
-              x, img = block_forward(block, x, img, styles2[:, styles_idx:styles_idx+3, :], temp_shapes[k], noise_mode=noise_mode)
-              styles_idx += 3
-
-        identity_loss, _ = id_loss(img, img2)
-        identity_loss *= identity_coefficient
-        img = (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255)
-        img = (transf(img.permute(0, 3, 1, 2))/255).sub_(mean).div_(std)
-        image_features = model.encode_image(img)
-        cos_sim = -1*F.cosine_similarity(image_features, (text_features[0]).unsqueeze(0))
-        (identity_loss + cos_sim.sum()).backward(retain_graph=True)
-
-
-
-
-    #t1 = time.time()
-
-    for text_counter in range(len(text_prompts)):
-        text_prompt = text_prompts[text_counter]
-        print(text_prompt)
-
-        styles_wanted_direction.grad.data.zero_()
-        styles_wanted_direction_grad_el2 = torch.zeros(1,26,512, device=device)
-        with torch.no_grad():
-            styles_wanted_direction *= 0
-
-        for i in range(math.ceil(len(seeds)/batch_size)):
-            print(i*batch_size, "processed", time.time()-t1)
-
-
-            styles = torch.vstack(styles_array[i*batch_size:(i+1)*batch_size]).to(device)
-
-
-            seed = seeds[i]
-
-            img2 = torch.tensor(temp_photos[i]).to(device)
-
-            styles2 = styles + styles_wanted_direction
-
-            styles_idx = 0
-            x = img = None
-            for k , (res, cur_ws) in enumerate(zip(G.synthesis.block_resolutions, block_ws)):
-                block = getattr(G.synthesis, f'b{res}')
-                if k>resolution_dict[resolution]:
-                  continue
-
-                if res == 4:
-                  x, img = block_forward(block, x, img, styles2[:, styles_idx:styles_idx+2, :], temp_shapes[k], noise_mode=noise_mode)
-                  styles_idx += 2
-                else:
-                  x, img = block_forward(block, x, img, styles2[:, styles_idx:styles_idx+3, :], temp_shapes[k], noise_mode=noise_mode)
-                  styles_idx += 3
-
-            identity_loss, _ = id_loss(img, img2)
-            identity_loss *= identity_coefficient
-            img = (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255)
-            img = (transf(img.permute(0, 3, 1, 2))/255).sub_(mean).div_(std)
-            image_features = model.encode_image(img)
-            cos_sim = -1*F.cosine_similarity(image_features, (text_features[text_counter]).unsqueeze(0))
-            (identity_loss + cos_sim.sum()).backward(retain_graph=True)
-
-
-            styles_wanted_direction.grad[:, [0, 1, 4, 7, 10, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25], :] = 0
-
-
-            if i%2==1:
-              styles_wanted_direction.data = styles_wanted_direction - styles_wanted_direction.grad*5
-              grads.append(styles_wanted_direction.grad.clone())
-              styles_wanted_direction.grad.data.zero_()
-
-              if i>3:
-                styles_wanted_direction_grad_el2[grads[-2]*grads[-1]<0] += 1
-
-
-        styles_wanted_direction_cpu = styles_wanted_direction.detach()
-        styles_wanted_direction_cpu[styles_wanted_direction_grad_el2>(len(seeds)/batch_size)/4] = 0
-        np.savez(f'{outdir}/direction_'+text_prompt.replace(" ", "_")+'.npz', s=styles_wanted_direction_cpu.cpu().numpy())
-
-    print("time passed:", time.time()-t1)
-#----------------------------------------------------------------------------
-
-if __name__ == "__main__":
-    generate_images() # pylint: disable=no-value-for-parameter
-
-#----------------------------------------------------------------------------

generate_w.py

@@ -1,148 +0,0 @@
-# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
-#
-# NVIDIA CORPORATION and its licensors retain all intellectual property
-# and proprietary rights in and to this software, related documentation
-# and any modifications thereto.  Any use, reproduction, disclosure or
-# distribution of this software and related documentation without an express
-# license agreement from NVIDIA CORPORATION is strictly prohibited.
-
-"""Generate images using pretrained network pickle."""
-
-import os
-import re
-from typing import List, Optional
-import torchvision
-from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize
-import click
-import dnnlib
-import numpy as np
-import PIL.Image
-import torch
-from torch import linalg as LA
-import clip
-from PIL import Image
-import legacy
-import torch.nn.functional as F
-import cv2
-import matplotlib.pyplot as plt
-from torch_utils import misc
-from torch_utils import persistence
-from torch_utils.ops import conv2d_resample
-from torch_utils.ops import upfirdn2d
-from torch_utils.ops import bias_act
-from torch_utils.ops import fma
-import random
-import math
-import time
-import id_loss
-
-
-def block_forward(self, x, img, ws, shapes, force_fp32=False, fused_modconv=None, **layer_kwargs):
-        misc.assert_shape(ws, [None, self.num_conv + self.num_torgb, self.w_dim])
-        w_iter = iter(ws.unbind(dim=1))
-        dtype = torch.float16 if self.use_fp16 and not force_fp32 else torch.float32
-        memory_format = torch.channels_last if self.channels_last and not force_fp32 else torch.contiguous_format
-        if fused_modconv is None:
-            with misc.suppress_tracer_warnings(): # this value will be treated as a constant
-                fused_modconv = (not self.training) and (dtype == torch.float32 or int(x.shape[0]) == 1)
-
-        # Input.
-        if self.in_channels == 0:
-            x = self.const.to(dtype=dtype, memory_format=memory_format)
-            x = x.unsqueeze(0).repeat([ws.shape[0], 1, 1, 1])
-        else:
-            misc.assert_shape(x, [None, self.in_channels, self.resolution // 2, self.resolution // 2])
-            x = x.to(dtype=dtype, memory_format=memory_format)
-
-        # Main layers.
-        if self.in_channels == 0:
-            x = self.conv1(x, next(w_iter)[...,:shapes[0]], fused_modconv=fused_modconv, **layer_kwargs)
-        elif self.architecture == 'resnet':
-            y = self.skip(x, gain=np.sqrt(0.5))
-            x = self.conv0(x, next(w_iter), fused_modconv=fused_modconv, **layer_kwargs)
-            x = self.conv1(x, next(w_iter), fused_modconv=fused_modconv, gain=np.sqrt(0.5), **layer_kwargs)
-            x = y.add_(x)
-        else:
-            x = self.conv0(x, next(w_iter)[...,:shapes[0]], fused_modconv=fused_modconv, **layer_kwargs)
-            x = self.conv1(x, next(w_iter)[...,:shapes[1]], fused_modconv=fused_modconv, **layer_kwargs)
-
-        # ToRGB.
-        if img is not None:
-            misc.assert_shape(img, [None, self.img_channels, self.resolution // 2, self.resolution // 2])
-            img = upfirdn2d.upsample2d(img, self.resample_filter)
-        if self.is_last or self.architecture == 'skip':
-            y = self.torgb(x, next(w_iter)[...,:shapes[2]], fused_modconv=fused_modconv)
-            y = y.to(dtype=torch.float32, memory_format=torch.contiguous_format)
-            img = img.add_(y) if img is not None else y
-
-        assert x.dtype == dtype
-        assert img is None or img.dtype == torch.float32
-        return x, img
-
-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]
-
-
-@click.command()
-@click.pass_context
-@click.option('--network', 'network_pkl', help='Network pickle filename', required=True)
-@click.option('--seeds', type=num_range, help='List of random seeds')
-@click.option('--trunc', 'truncation_psi', type=float, help='Truncation psi', default=1, show_default=True)
-@click.option('--class', 'class_idx', type=int, help='Class label (unconditional if not specified)')
-@click.option('--noise-mode', help='Noise mode', type=click.Choice(['const', 'random', 'none']), default='const', show_default=True)
-def generate_images(
-    ctx: click.Context,
-    network_pkl: str,
-    seeds: Optional[List[int]],
-    truncation_psi: float,
-    noise_mode: str,
-    class_idx: Optional[int],
-    projected_w: Optional[str],
-    projected_s: Optional[str]
-):
-
-    print('Loading networks from "%s"...' % network_pkl)
-    # Use GPU if available
-    if torch.cuda.is_available():
-        device = torch.device("cuda")
-    else:
-        device = torch.device("cpu")
-
-    with dnnlib.util.open_url(network_pkl) as f:
-        G = legacy.load_network_pkl(f)['G_ema'].to(device) # type: ignore
-
-    if seeds is None:
-        ctx.fail('--seeds option is required when not using --projected-w')
-
-    # Labels.
-    label = torch.zeros([1, G.c_dim], device=device).requires_grad_()
-    if G.c_dim != 0:
-        if class_idx is None:
-            ctx.fail('Must specify class label with --class when using a conditional network')
-        label[:, class_idx] = 1
-    else:
-        if class_idx is not None:
-            print ('warn: --class=lbl ignored when running on an unconditional network')
-
-    z = torch.from_numpy(np.random.RandomState(seeds[0]).randn(1, G.z_dim)).to(device)
-    ws = G.mapping(z, label, truncation_psi=truncation_psi)
-    np.savez(f'encoder4editing/projected_w.npz', w=ws.detach().cpu().numpy())
-
-
-if __name__ == "__main__":
-    generate_images()

w_s_converter.py

@@ -10,7 +10,7 @@
 
 import os
 import re
-from typing import List
+from typing import List, Optional
 
 import numpy as np
 import torch
@@ -23,6 +23,13 @@ from torch_utils.ops import bias_act
 from torch_utils.ops import fma
 
 
+import click
+
+import PIL.Image
+from torch import linalg as LA
+import torch.nn.functional as F
+
+
 def block_forward(self, x, img, ws, shapes, force_fp32=False, fused_modconv=None, **layer_kwargs):
         misc.assert_shape(ws, [None, self.num_conv + self.num_torgb, self.w_dim])
         w_iter = iter(ws.unbind(dim=1))
@@ -66,13 +73,56 @@ def block_forward(self, x, img, ws, shapes, force_fp32=False, fused_modconv=None
         return x, img
 
 
+def block_forward_from_style(self, x, img, ws, shapes, force_fp32=False, fused_modconv=None, **layer_kwargs):
+        misc.assert_shape(ws, [None, self.num_conv + self.num_torgb, self.w_dim])
+        w_iter = iter(ws.unbind(dim=1))
+        dtype = torch.float16 if self.use_fp16 and not force_fp32 else torch.float32
+        memory_format = torch.channels_last if self.channels_last and not force_fp32 else torch.contiguous_format
+        if fused_modconv is None:
+            with misc.suppress_tracer_warnings(): # this value will be treated as a constant
+                fused_modconv = (not self.training) and (dtype == torch.float32 or int(x.shape[0]) == 1)
+
+        # Input.
+        if self.in_channels == 0:
+            x = self.const.to(dtype=dtype, memory_format=memory_format)
+            x = x.unsqueeze(0).repeat([ws.shape[0], 1, 1, 1])
+        else:
+            misc.assert_shape(x, [None, self.in_channels, self.resolution // 2, self.resolution // 2])
+            x = x.to(dtype=dtype, memory_format=memory_format)
+
+        # Main layers.
+        if self.in_channels == 0:
+            x = self.conv1(x, next(w_iter)[...,:shapes[0]], fused_modconv=fused_modconv, **layer_kwargs)
+        elif self.architecture == 'resnet':
+            y = self.skip(x, gain=np.sqrt(0.5))
+            x = self.conv0(x, next(w_iter), fused_modconv=fused_modconv, **layer_kwargs)
+            x = self.conv1(x, next(w_iter), fused_modconv=fused_modconv, gain=np.sqrt(0.5), **layer_kwargs)
+            x = y.add_(x)
+        else:
+            x = self.conv0(x, next(w_iter)[...,:shapes[0]], fused_modconv=fused_modconv, **layer_kwargs)
+            x = self.conv1(x, next(w_iter)[...,:shapes[1]], fused_modconv=fused_modconv, **layer_kwargs)
+
+        # ToRGB.
+        if img is not None:
+            misc.assert_shape(img, [None, self.img_channels, self.resolution // 2, self.resolution // 2])
+            img = upfirdn2d.upsample2d(img, self.resample_filter)
+        if self.is_last or self.architecture == 'skip':
+            y = self.torgb(x, next(w_iter)[...,:shapes[2]], fused_modconv=fused_modconv)
+            y = y.to(dtype=torch.float32, memory_format=torch.contiguous_format)
+            img = img.add_(y) if img is not None else y
+
+        assert x.dtype == dtype
+        assert img is None or img.dtype == torch.float32
+        return x, img
+
+
 def unravel_index(index, shape):
     out = []
     for dim in reversed(shape):
         out.append(index % dim)
         index = index // dim
     return tuple(reversed(out))
-    
+
 
 def w_to_s(
     G,
@@ -136,3 +186,75 @@ def w_to_s(
 
     styles = styles.detach()
     np.savez(f'{outdir}/input.npz', s=styles.cpu().numpy())
+
+
+def generate_from_style(
+    G,
+    outdir: str,
+    s_input: str,
+    text_prompt: str,
+    change_power: int,
+    truncation_psi: float = 0.7,
+    noise_mode: str = "const",
+):
+    # Use GPU if available
+    if torch.cuda.is_available():
+        device = torch.device("cuda")
+    else:
+        device = torch.device("cpu")
+
+    os.makedirs(outdir, exist_ok=True)
+
+    # Generate images
+    for i in G.parameters():
+      i.requires_grad = False
+
+    temp_shapes = []
+    for res in G.synthesis.block_resolutions:
+        block = getattr(G.synthesis, f'b{res}')
+        if res == 4:
+            temp_shape = (block.conv1.affine.weight.shape[0], block.conv1.affine.weight.shape[0], block.torgb.affine.weight.shape[0])
+            block.conv1.affine = torch.nn.Identity()
+            block.torgb.affine = torch.nn.Identity()
+
+        else:
+            temp_shape = (block.conv0.affine.weight.shape[0], block.conv1.affine.weight.shape[0], block.torgb.affine.weight.shape[0])
+            block.conv0.affine = torch.nn.Identity()
+            block.conv1.affine = torch.nn.Identity()
+            block.torgb.affine = torch.nn.Identity()
+
+        temp_shapes.append(temp_shape)
+
+    if s_input is not None:
+        styles = np.load(s_input)['s']
+        styles_direction = np.load(f'{outdir}/direction_'+text_prompt.replace(" ", "_")+'.npz')['s']
+
+        styles_direction = torch.tensor(styles_direction, device=device)
+        styles = torch.tensor(styles, device=device)
+
+    with torch.no_grad():
+        imgs = []
+        grad_changes = [0, 0.25*change_power, 0.5*change_power, 0.75*change_power, change_power]
+
+        for grad_change in grad_changes:
+            styles += styles_direction*grad_change
+
+            styles_idx = 0
+            x = img = None
+            for k , res in enumerate(G.synthesis.block_resolutions):
+                block = getattr(G.synthesis, f'b{res}')
+
+                if res == 4:
+                    x, img = block_forward_from_style(block, x, img, styles[:, styles_idx:styles_idx+2, :], temp_shapes[k], noise_mode=noise_mode, force_fp32=True)
+                    styles_idx += 2
+                else:
+                    x, img = block_forward_from_style(block, x, img, styles[:, styles_idx:styles_idx+3, :], temp_shapes[k], noise_mode=noise_mode, force_fp32=True)
+                    styles_idx += 3
+
+            img = (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255)
+            imgs.append(img[0].to(torch.uint8).cpu().numpy())
+
+            styles -= styles_direction*grad_change
+
+        img_filepath = f'{outdir}/'+text_prompt.replace(" ", "_")+'_'+str(change_power)+'.jpeg'
+        PIL.Image.fromarray(np.concatenate(imgs, axis=1), 'RGB').save(img_filepath, quality=95)