Allow user to generate images with initial noise as on M1 / mps system

ldm/dream/args.py

@@ -753,6 +753,13 @@ def _create_dream_cmd_parser(self):
             type=str,
             help='list of variations to apply, in the format `seed:weight,seed:weight,...'
         )
+        render_group.add_argument(
+            '--use_mps_noise',
+            action='store_true',
+            dest='use_mps_noise',
+            help='Simulate noise on M1 systems to get the same results'
+        )
+
         return parser
 
 def format_metadata(**kwargs):

ldm/dream/generator/base.py

@@ -23,6 +23,7 @@ def __init__(self, model, precision):
         self.downsampling_factor = downsampling   # BUG: should come from model or config
         self.variation_amount    = 0
         self.with_variations     = []
+        self.use_mps_noise       = False
 
     # this is going to be overridden in img2img.py, txt2img.py and inpaint.py
     def get_make_image(self,prompt,**kwargs):

ldm/dream/generator/txt2img.py

@@ -59,7 +59,7 @@ def make_image(x_T):
     # returns a tensor filled with random numbers from a normal distribution
     def get_noise(self,width,height):
         device         = self.model.device
-        if device.type == 'mps':
+        if self.use_mps_noise or device.type == 'mps':
             x = torch.randn([1,
                                 self.latent_channels,
                                 height // self.downsampling_factor,

ldm/dream/generator/txt2img2img.py

@@ -116,7 +116,7 @@ def get_noise(self,width,height,scale = True):
             scaled_height = height
 
         device      = self.model.device
-        if device.type == 'mps':
+        if self.use_mps_noise or device.type == 'mps':
             return torch.randn([1,
                                 self.latent_channels,
                                 scaled_height // self.downsampling_factor,

ldm/generate.py

@@ -258,6 +258,7 @@ def prompt2image(
             # Set this True to handle KeyboardInterrupt internally
             catch_interrupts = False,
             hires_fix        = False,
+            use_mps_noise    = False,
             **args,
     ):   # eat up additional cruft
         """
@@ -386,6 +387,7 @@ def process_image(image,seed):
 
             generator.set_variation(
                 self.seed, variation_amount, with_variations)
+            generator.use_mps_noise = use_mps_noise
             results = generator.generate(
                 prompt,
                 iterations=iterations,