fixes dream.py mps seed

lstein · lstein · commit fe5cc792495e · 2022-09-03T10:11:46.000-04:00
diff --git a/ldm/simplet2i.py b/ldm/simplet2i.py
@@ -348,13 +348,23 @@ def process_image(image,seed):
 
             def get_noise():
                 if init_img:
-                    return torch.randn_like(init_latent, device=self.device)
+                    if self.device.type == 'mps':
+                        return torch.randn_like(init_latent, device='cpu').to(self.device)
+                    else:
+                        return torch.randn_like(init_latent, device=self.device)
                 else:
-                    return torch.randn([1,
-                        self.latent_channels,
-                        height // self.downsampling_factor,
-                        width  // self.downsampling_factor],
-                        device=self.device)
+                    if self.device.type == 'mps':
+                        return torch.randn([1,
+                            self.latent_channels,
+                            height // self.downsampling_factor,
+                            width  // self.downsampling_factor],
+                            device='cpu').to(self.device)
+                    else:
+                        return torch.randn([1,
+                            self.latent_channels,
+                            height // self.downsampling_factor,
+                            width  // self.downsampling_factor],
+                            device=self.device)
 
             initial_noise = None
             if variation_amount > 0 or len(with_variations) > 0:
@@ -383,6 +393,8 @@ def get_noise():
                         x_T = initial_noise
                     else:
                         seed_everything(seed)
+                        if self.device.type == 'mps':
+                            x_T = get_noise()
                         # make_image will do the equivalent of get_noise itself
                     image = make_image(x_T)
                     results.append([image, seed])
