no styleclip

app.py

@@ -18,8 +18,7 @@ import torchvision.transforms as transforms
 from torchvision import utils
 
 from model.sg2_model import Generator
-from generate_videos import generate_frames, video_from_interpolations, project_code_by_edit_name
-from styleclip.styleclip_global import project_code_with_styleclip, style_tensor_to_style_dict
+from generate_videos import project_code_by_edit_name
 
 import clip
 
@@ -30,7 +29,6 @@ os.makedirs(model_dir, exist_ok=True)
 model_repos = {
     "e4e": ("akhaliq/JoJoGAN_e4e_ffhq_encode", "e4e_ffhq_encode.pt"),
                "dlib": ("akhaliq/jojogan_dlib", "shape_predictor_68_face_landmarks.dat"),
-               "sc_fs3": ("rinong/stylegan-nada-models", "fs3.npy"),
                "base": ("akhaliq/jojogan-stylegan2-ffhq-config-f", "stylegan2-ffhq-config-f.pt"),
                "sketch": ("rinong/stylegan-nada-models", "sketch.pt"),
                 "santa": ("mjdolan/stylegan-nada-models", "santa.pt"),
@@ -78,7 +76,7 @@ class ImageEditor(object):
 
         self.generators = {}
 
-        self.model_list = [name for name in model_paths.keys() if name not in ["e4e", "dlib", "sc_fs3"]]
+        self.model_list = [name for name in model_paths.keys() if name not in ["e4e", "dlib"]]
 
         for model in self.model_list:
             g_ema = Generator(
@@ -116,7 +114,6 @@ class ImageEditor(object):
             model_paths["dlib"]
         )
 
-        self.styleclip_fs3 = torch.from_numpy(np.load(model_paths["sc_fs3"])).to(self.device)
 
         self.clip_model, _ = clip.load("ViT-B/32", device=self.device)
 
@@ -163,6 +160,7 @@ class ImageEditor(object):
 
         return inner
 
+
     def get_target_latent(self, source_latent, alter, generators):
         np_source_latent = source_latent.squeeze(0).cpu().detach().numpy()
         if alter == "None":

styleclip/styleclip_global.py

@@ -1,181 +0,0 @@
-'''
-Code adapted from Stitch it in Time by Tzaban et al.
-https://github.com/rotemtzaban/STIT 
-'''
-
-
-import numpy as np
-import torch
-from tqdm import tqdm
-from pathlib import Path
-import os
-
-import clip
-
-imagenet_templates = [
-    'a bad photo of a {}.',
-    'a photo of many {}.',
-    'a sculpture of a {}.',
-    'a photo of the hard to see {}.',
-    'a low resolution photo of the {}.',
-    'a rendering of a {}.',
-    'graffiti of a {}.',
-    'a bad photo of the {}.',
-    'a cropped photo of the {}.',
-    'a tattoo of a {}.',
-    'the embroidered {}.',
-    'a photo of a hard to see {}.',
-    'a bright photo of a {}.',
-    'a photo of a clean {}.',
-    'a photo of a dirty {}.',
-    'a dark photo of the {}.',
-    'a drawing of a {}.',
-    'a photo of my {}.',
-    'the plastic {}.',
-    'a photo of the cool {}.',
-    'a close-up photo of a {}.',
-    'a black and white photo of the {}.',
-    'a painting of the {}.',
-    'a painting of a {}.',
-    'a pixelated photo of the {}.',
-    'a sculpture of the {}.',
-    'a bright photo of the {}.',
-    'a cropped photo of a {}.',
-    'a plastic {}.',
-    'a photo of the dirty {}.',
-    'a jpeg corrupted photo of a {}.',
-    'a blurry photo of the {}.',
-    'a photo of the {}.',
-    'a good photo of the {}.',
-    'a rendering of the {}.',
-    'a {} in a video game.',
-    'a photo of one {}.',
-    'a doodle of a {}.',
-    'a close-up photo of the {}.',
-    'a photo of a {}.',
-    'the origami {}.',
-    'the {} in a video game.',
-    'a sketch of a {}.',
-    'a doodle of the {}.',
-    'a origami {}.',
-    'a low resolution photo of a {}.',
-    'the toy {}.',
-    'a rendition of the {}.',
-    'a photo of the clean {}.',
-    'a photo of a large {}.',
-    'a rendition of a {}.',
-    'a photo of a nice {}.',
-    'a photo of a weird {}.',
-    'a blurry photo of a {}.',
-    'a cartoon {}.',
-    'art of a {}.',
-    'a sketch of the {}.',
-    'a embroidered {}.',
-    'a pixelated photo of a {}.',
-    'itap of the {}.',
-    'a jpeg corrupted photo of the {}.',
-    'a good photo of a {}.',
-    'a plushie {}.',
-    'a photo of the nice {}.',
-    'a photo of the small {}.',
-    'a photo of the weird {}.',
-    'the cartoon {}.',
-    'art of the {}.',
-    'a drawing of the {}.',
-    'a photo of the large {}.',
-    'a black and white photo of a {}.',
-    'the plushie {}.',
-    'a dark photo of a {}.',
-    'itap of a {}.',
-    'graffiti of the {}.',
-    'a toy {}.',
-    'itap of my {}.',
-    'a photo of a cool {}.',
-    'a photo of a small {}.',
-    'a tattoo of the {}.',
-]
-
-CONV_CODE_INDICES = [(0, 512), (1024, 1536), (1536, 2048), (2560, 3072), (3072, 3584), (4096, 4608), (4608, 5120), (5632, 6144), (6144, 6656), (7168, 7680), (7680, 7936), (8192, 8448), (8448, 8576), (8704, 8832), (8832, 8896), (8960, 9024), (9024, 9056)]
-FFHQ_CODE_INDICES = [(0, 512), (512, 1024), (1024, 1536), (1536, 2048), (2560, 3072), (3072, 3584), (4096, 4608), (4608, 5120), (5632, 6144), (6144, 6656), (7168, 7680), (7680, 7936), (8192, 8448), (8448, 8576), (8704, 8832), (8832, 8896), (8960, 9024), (9024, 9056)] + \
-                    [(2048, 2560), (3584, 4096), (5120, 5632), (6656, 7168), (7936, 8192), (8576, 8704), (8896, 8960), (9056, 9088)]
-
-def zeroshot_classifier(model, classnames, templates, device):
-
-    with torch.no_grad():
-        zeroshot_weights = []
-        for classname in tqdm(classnames):
-            texts = [template.format(classname) for template in templates]  # format with class
-            texts = clip.tokenize(texts).to(device)  # tokenize
-            class_embeddings = model.encode_text(texts)  # embed with text encoder
-            class_embeddings /= class_embeddings.norm(dim=-1, keepdim=True)
-            class_embedding = class_embeddings.mean(dim=0)
-            class_embedding /= class_embedding.norm()
-            zeroshot_weights.append(class_embedding)
-        zeroshot_weights = torch.stack(zeroshot_weights, dim=1).to(device)
-    return zeroshot_weights
-
-def expand_to_full_dim(partial_tensor):
-    full_dim_tensor = torch.zeros(size=(1, 9088))
-
-    start_idx = 0
-    for conv_start, conv_end in CONV_CODE_INDICES:
-        length = conv_end - conv_start
-        full_dim_tensor[:, conv_start:conv_end] = partial_tensor[start_idx:start_idx + length]
-        start_idx += length
-
-    return full_dim_tensor
-
-def get_direction(neutral_class, target_class, beta, di, clip_model=None):
-
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-
-    if clip_model is None:
-        clip_model, _ = clip.load("ViT-B/32", device=device)
-
-    class_names = [neutral_class, target_class]
-    class_weights = zeroshot_classifier(clip_model, class_names, imagenet_templates, device)
-
-    dt = class_weights[:, 1] - class_weights[:, 0]
-    dt = dt / dt.norm()
-
-    dt = dt.float()
-    di = di.float()
-
-    relevance = di @ dt
-    mask = relevance.abs() > beta
-    direction = relevance * mask
-    direction_max = direction.abs().max()
-    if direction_max > 0:
-        direction = direction / direction_max
-    else:
-        raise ValueError(f'Beta value {beta} is too high for mapping from {neutral_class} to {target_class},'
-                         f' try setting it to a lower value')
-    return direction
-
-def style_tensor_to_style_dict(style_tensor, refernce_generator):
-    style_layers = refernce_generator.modulation_layers
-
-    style_dict = {}
-    for layer_idx, layer in enumerate(style_layers):
-        style_dict[layer] = style_tensor[:, FFHQ_CODE_INDICES[layer_idx][0]:FFHQ_CODE_INDICES[layer_idx][1]]
-
-    return style_dict
-
-def style_dict_to_style_tensor(style_dict, reference_generator):
-    style_layers = reference_generator.modulation_layers
-
-    style_tensor = torch.zeros(size=(1, 9088))
-    for layer in style_dict:
-        layer_idx = style_layers.index(layer)
-        style_tensor[:, FFHQ_CODE_INDICES[layer_idx][0]:FFHQ_CODE_INDICES[layer_idx][1]] = style_dict[layer]
-
-    return style_tensor
-
-def project_code_with_styleclip(source_latent, source_class, target_class, alpha, beta, reference_generator, di, clip_model=None):
-        edit_direction = get_direction(source_class, target_class, beta, di, clip_model)
-
-        edit_full_dim = expand_to_full_dim(edit_direction)
-
-        source_s = style_dict_to_style_tensor(source_latent, reference_generator)
-
-        return source_s + alpha * edit_full_dim