update find_direction

find_direction.py

@@ -8,34 +8,17 @@
 
 """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 torch
-from torch import linalg as LA
 import torch.nn.functional as F
-import torchvision
-from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize
-import PIL.Image
+from torchvision.transforms import Compose, Resize, CenterCrop
 from PIL import Image
-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 id_loss
 
 
@@ -81,8 +64,6 @@ def block_forward(self, x, img, ws, shapes, force_fp32=False, fused_modconv=None
         assert img is None or img.dtype == torch.float32
         return x, img
 
-
-
 def unravel_index(index, shape):
     out = []
     for dim in reversed(shape):
@@ -90,108 +71,27 @@ def unravel_index(index, shape):
         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('--text_prompt', help='Text', type=str, required=True)
-@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,
+def find_direction(
     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],
     text_prompt: str,
-    resolution: int,
-    batch_size: int,
-    identity_power: str,
+    truncation_psi: float = 0.7,
+    noise_mode: str = "const",
+    resolution: int = 256,
+    identity_power: float = 0.5,
 ):
-    """
-    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
-    """
+    seeds=np.random.randint(0, 1000, 128)
 
+    batch_size=1
     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")
-
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "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').save(f'{outdir}/proj{idx:02d}.png')
-        return
-
-    if seeds is None:
-        ctx.fail('--seeds option is required when not using --projected-w')
-
     # Labels
+    class_idx=None
     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 = clip.tokenize([text_prompt]).to(device)
@@ -211,8 +111,6 @@ def generate_images(
     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)))
@@ -260,8 +158,7 @@ def generate_images(
         styles_array.append(styles)
 
     resolution_dict = {256: 6, 512: 7, 1024: 8}
-    id_coeff_dict = {"high": 2, "medium": 0.5, "low": 0.1, "none": 0}
-    id_coeff = id_coeff_dict[identity_power]
+    id_coeff = identity_power
     styles_direction = torch.zeros(1, 26, 512, device=device)
     styles_direction_grad_el2 = torch.zeros(1, 26, 512, device=device)
     styles_direction.requires_grad_()
@@ -272,7 +169,6 @@ def generate_images(
     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]
@@ -325,6 +221,7 @@ def generate_images(
         styles_direction *= 0
 
     for i in range(math.ceil(len(seeds) / batch_size)):
+
         seed = seeds[i]
         styles = torch.vstack(styles_array[i*batch_size:(i+1)*batch_size]).to(device)
         img2 = torch.tensor(temp_photos[i]).to(device)
@@ -364,9 +261,4 @@ def generate_images(
     styles_direction = styles_direction.detach()
     styles_direction[styles_direction_grad_el2 > (len(seeds) / batch_size) / 4] = 0
 
-    output_filepath = f'{outdir}/direction_' + text_prompt.replace(" ", "_") + '.npz'
-    np.savez(output_filepath, s=styles_direction.cpu().numpy())
-
-
-if __name__ == "__main__":
-    generate_images()
+    return styles_direction.cpu().numpy()