add clip and opencv to the requirements

.gitignore

@@ -1,5 +1,6 @@
 # Custom
 .DS_Store
+.idea/
 
 # Byte-compiled / optimized / DLL files
 __pycache__/

kandinsky2/kandinsky2_2_model.py

@@ -21,12 +21,18 @@ def __init__(
     ):
         self.device = device
         self.task_type = task_type
-        self.image_encoder = CLIPVisionModelWithProjection.from_pretrained('kandinsky-community/kandinsky-2-2-prior', subfolder='image_encoder').to(torch.float16).to(self.device)
+
+        if device == "cpu":
+            torch_dtype = torch.float32
+        else:
+            torch_dtype = torch.float16
+
+        self.image_encoder = CLIPVisionModelWithProjection.from_pretrained('kandinsky-community/kandinsky-2-2-prior', subfolder='image_encoder').to(torch_dtype).to(self.device)
         if task_type == "text2img":
-            self.unet = UNet2DConditionModel.from_pretrained('kandinsky-community/kandinsky-2-2-decoder', subfolder='unet').to(torch.float16).to(self.device)
-            self.prior = KandinskyV22PriorPipeline.from_pretrained('kandinsky-community/kandinsky-2-2-prior', image_encoder=self.image_encoder, torch_dtype=torch.float16)
+            self.unet = UNet2DConditionModel.from_pretrained('kandinsky-community/kandinsky-2-2-decoder', subfolder='unet').to(torch_dtype).to(self.device)
+            self.prior = KandinskyV22PriorPipeline.from_pretrained('kandinsky-community/kandinsky-2-2-prior', image_encoder=self.image_encoder, torch_dtype=torch_dtype)
             self.prior = self.prior.to(self.device)
-            self.decoder = KandinskyV22Pipeline.from_pretrained('kandinsky-community/kandinsky-2-2-decoder', unet=self.unet, torch_dtype=torch.float16)
+            self.decoder = KandinskyV22Pipeline.from_pretrained('kandinsky-community/kandinsky-2-2-decoder', unet=self.unet, torch_dtype=torch_dtype)
             self.decoder = self.decoder.to(self.device)
         elif task_type == "inpainting":
             self.unet = UNet2DConditionModel.from_pretrained('kandinsky-community/kandinsky-2-2-decoder-inpaint', subfolder='unet').to(torch.float16).to(self.device)

kandinsky2/model/gaussian_diffusion.py

@@ -6,6 +6,7 @@
 
 import enum
 import math
+import traceback
 from copy import deepcopy
 import numpy as np
 import torch as th
@@ -822,7 +823,12 @@ def _extract_into_tensor(arr, timesteps, broadcast_shape):
                             dimension equal to the length of timesteps.
     :return: a tensor of shape [batch_size, 1, ...] where the shape has K dims.
     """
-    res = th.from_numpy(arr).to(device=timesteps.device)[timesteps].float()
+    try:
+        # For CUDA
+        res = th.from_numpy(arr).to(device=timesteps.device)[timesteps].float()
+    except Exception as e:
+        # For MPS
+        res = th.from_numpy(arr).to(device=timesteps.device, dtype=th.float32)[timesteps].float()
     while len(res.shape) < len(broadcast_shape):
         res = res[..., None]
     return res.expand(broadcast_shape)

setup.py

@@ -25,8 +25,9 @@
         "einops",
         "sentencepiece",
         "diffusers",
-        "accelerate"
-
+        "accelerate",
+        "clip @ git+https://github.com/openai/CLIP.git",
+        "opencv-python"
     ],
     author="",
 )

train_2_1_unclip.py

@@ -20,6 +20,7 @@
 import clip
 import argparse
 
+
 def drop_first_layer(path):
     d = {}
     state_dict = torch.load(path)
@@ -28,6 +29,7 @@ def drop_first_layer(path):
             d[key] = state_dict[key]
     return d
 
+
 def main():
     parser = argparse.ArgumentParser()
     parser.add_argument('--config', type=str, help='config path')
@@ -65,11 +67,13 @@ def main():
     clip_model.token_embedding = None
     clip_model.text_projection = None
     clip_model = clip_model.eval().to(device)
-    train_unclip(unet=model, diffusion=diffusion, image_encoder=image_encoder,
-                  clip_model=clip_model, text_encoder=text_encoder, optimizer=optimizer,
-                  lr_scheduler=lr_scheduler, schedule_sampler=schedule_sampler, 
-                  train_loader=train_loader, val_loader=None, scale=config['image_enc_params']['scale'],
-                  num_epochs=config['num_epochs'], save_every=config['save_every'], save_name=config['save_name'],
-                  save_path=config['save_path'],  inpainting=config['inpainting'], device=device)
+    train_unclip(unet=model, diffusion=diffusion, image_encoder=image_encoder, clip_model=clip_model,
+                 text_encoder=text_encoder, optimizer=optimizer, lr_scheduler=lr_scheduler,
+                 schedule_sampler=schedule_sampler, train_loader=train_loader, val_loader=None,
+                 scale=config['image_enc_params']['scale'], num_epochs=config['num_epochs'],
+                 save_every=config['save_every'], save_name=config['save_name'], save_path=config['save_path'],
+                 inpainting=config['inpainting'], device=device)
+
+
 if __name__ == '__main__':
     main()