Update stable_diffusion.py

wangcaihua · Oct 30, 2022 · 41973c2 · 41973c2
1 parent b5b0fa5
commit 41973c2
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/stable_diffusion_tf/stable_diffusion.py b/stable_diffusion_tf/stable_diffusion.py
@@ -68,7 +68,7 @@ def generate(
                 input_image = input_image.resize((self.img_width, self.img_height))
 
             elif type(input_image) is np.ndarray:
-                input_image = np.resize(input_image, (self.img_width, self.img_height, input_image.shape[2]))
+                input_image = np.resize(input_image, (self.img_height, self.img_width, input_image.shape[2]))
 
             input_image_array = np.array(input_image, dtype=np.float32)[None,...,:3]
             input_image_tensor = tf.cast((input_image_array / 255.0) * 2 - 1, self.dtype)
@@ -200,7 +200,7 @@ def get_x_prev_and_pred_x0(self, x, e_t, index, a_t, a_prev, temperature, seed):
 
     def load_weights_from_pytorch_ckpt(self , pytorch_ckpt_path):
         import torch
-        pt_weights = torch.load(pytorch_ckpt_path)
+        pt_weights = torch.load(pytorch_ckpt_path, map_location="cpu")
         for module_name in ['text_encoder', 'diffusion_model', 'decoder', 'encoder' ]:
             module_weights = []
             for i , (key , perm ) in enumerate(PYTORCH_CKPT_MAPPING[module_name]):