Fixes to run on CPU and MPS

Author: Wojtek Kowaluk

kandinsky2/configs.py

@@ -136,7 +136,7 @@
         "model_dim": 768,
         "use_scale_shift_norm": True,
         "resblock_updown": True,
-        "use_fp16": True,
+        "use_fp16": False,
         "cache_text_emb": True,
         "text_encoder_in_dim1": 1024,
         "text_encoder_in_dim2": 768,

kandinsky2/kandinsky2_1_model.py

@@ -54,7 +54,7 @@ def __init__(
             clip_mean,
             clip_std,
         )
-        self.prior.load_state_dict(torch.load(prior_path), strict=False)
+        self.prior.load_state_dict(torch.load(prior_path, map_location='cpu'), strict=False)
         if self.use_fp16:
             self.prior = self.prior.half()
         self.text_encoder = TextEncoder(**self.config["text_enc_params"])
@@ -88,7 +88,7 @@ def __init__(
 
         self.config["model_config"]["cache_text_emb"] = True
         self.model = create_model(**self.config["model_config"])
-        self.model.load_state_dict(torch.load(model_path))
+        self.model.load_state_dict(torch.load(model_path, map_location='cpu'))
         if self.use_fp16:
             self.model.convert_to_fp16()
             self.image_encoder = self.image_encoder.half()

kandinsky2/model/gaussian_diffusion.py

@@ -822,7 +822,7 @@ 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()
+    res = th.from_numpy(arr).to(dtype=th.float32).to(device=timesteps.device)[timesteps]
     while len(res.shape) < len(broadcast_shape):
         res = res[..., None]
     return res.expand(broadcast_shape)