update models

clip_encoder.py

@@ -62,3 +62,24 @@ class CLIPImageEncoder(nn.Module):
 
 
 
+class OpenCLIPImageEncoder(nn.Module):
+
+    def __init__(self, model="ViT-B/32", pretrained="openai"):
+        super().__init__()
+        model, _, preprocess = open_clip.create_model_and_transforms(model, pretrained=pretrained)
+        self.tokenizer = open_clip.get_tokenizer(model)
+        CLIP_MEAN = [0.48145466, 0.4578275, 0.40821073]
+        CLIP_STD = [0.26862954, 0.26130258, 0.27577711]
+        mean = torch.tensor(CLIP_MEAN).view(1, 3, 1, 1)
+        std = torch.tensor(CLIP_STD).view(1, 3, 1, 1)
+        self.register_buffer("mean", mean)
+        self.register_buffer("std", std)
+
+    def forward_image(self, x):
+        x = torch.nn.functional.interpolate(x, mode='bicubic', size=(224, 224))
+        x = (x-self.mean)/self.std
+        return self.model.encode_image(x)
+
+    def forward_text(self, texts):
+        toks = self.tokenizer.tokenize(texts, truncate=True).to(self.mean.device)
+        return self.model.encode_text(toks)

run.py

@@ -237,7 +237,7 @@ def ddgan_laion2b_v2():
     return cfg
 
 def ddgan_ddb_v1():
-    cfg = ddgan_sd_v9()
+    cfg = ddgan_sd_v10()
     return cfg
 
 def ddgan_sd_v11():
@@ -245,6 +245,17 @@ def ddgan_sd_v11():
     cfg['model']['image_size'] = 512
     return cfg
 
+def ddgan_ddb_v2():
+    cfg = ddgan_ddb_v1()
+    cfg['model']['num_timesteps'] = 1
+    return cfg
+
+def ddgan_ddb_v3():
+    cfg = ddgan_ddb_v1()
+    cfg['model']['num_channels_dae'] = 192
+    cfg['model']['num_timesteps'] = 2
+    return cfg
+
 models = [
     ddgan_cifar10_cond17, # cifar10, cross attn for discr
     ddgan_cifar10_cond18, # cifar10, xl encoder
@@ -286,6 +297,8 @@ models = [
     ddgan_sd_v11,
     ddgan_laion2b_v2,
     ddgan_ddb_v1,
+    ddgan_ddb_v2,
+    ddgan_ddb_v3
 ]
 
 def get_model(model_name):