Fix generator determinism: forward generator through scheduler steps and seeded noise

edm2-img512-l-dino/pipeline.py

@@ -350,6 +350,9 @@ class EDM2Pipeline(DiffusionPipeline):
 
         device = self._execution_device
         dtype = self.unet.dtype
+        self.unet.eval()
+        if getattr(self, "gnet", None) is not None:
+            self.gnet.eval()
         labels = self._normalize_class_labels(class_labels, batch_size=batch_size, device=device)
         noise = self.prepare_latents(batch_size, height, width, dtype, device, generator)
 
@@ -375,13 +378,14 @@ class EDM2Pipeline(DiffusionPipeline):
         latents = self._sample_edm2_heun(
             denoise_fn=denoise_fn,
             noise=noise,
-            sigmas=self.scheduler.sigmas.to(device),
+            sigmas=self.scheduler.sigmas.to(device).clone(),
             generator=generator,
             progress_bar=self.progress_bar,
             dtype=torch.float32,
         )
 
         image = self.decode_latents(latents, output_type=output_type)
+        self.maybe_free_model_hooks()
         if not return_dict:
             return (image, latents)
         return ImagePipelineOutput(images=image)

edm2-img512-l-fid/pipeline.py

@@ -350,6 +350,9 @@ class EDM2Pipeline(DiffusionPipeline):
 
         device = self._execution_device
         dtype = self.unet.dtype
+        self.unet.eval()
+        if getattr(self, "gnet", None) is not None:
+            self.gnet.eval()
         labels = self._normalize_class_labels(class_labels, batch_size=batch_size, device=device)
         noise = self.prepare_latents(batch_size, height, width, dtype, device, generator)
 
@@ -375,13 +378,14 @@ class EDM2Pipeline(DiffusionPipeline):
         latents = self._sample_edm2_heun(
             denoise_fn=denoise_fn,
             noise=noise,
-            sigmas=self.scheduler.sigmas.to(device),
+            sigmas=self.scheduler.sigmas.to(device).clone(),
             generator=generator,
             progress_bar=self.progress_bar,
             dtype=torch.float32,
         )
 
         image = self.decode_latents(latents, output_type=output_type)
+        self.maybe_free_model_hooks()
         if not return_dict:
             return (image, latents)
         return ImagePipelineOutput(images=image)

edm2-img512-m-fid/pipeline.py

@@ -350,6 +350,9 @@ class EDM2Pipeline(DiffusionPipeline):
 
         device = self._execution_device
         dtype = self.unet.dtype
+        self.unet.eval()
+        if getattr(self, "gnet", None) is not None:
+            self.gnet.eval()
         labels = self._normalize_class_labels(class_labels, batch_size=batch_size, device=device)
         noise = self.prepare_latents(batch_size, height, width, dtype, device, generator)
 
@@ -375,13 +378,14 @@ class EDM2Pipeline(DiffusionPipeline):
         latents = self._sample_edm2_heun(
             denoise_fn=denoise_fn,
             noise=noise,
-            sigmas=self.scheduler.sigmas.to(device),
+            sigmas=self.scheduler.sigmas.to(device).clone(),
             generator=generator,
             progress_bar=self.progress_bar,
             dtype=torch.float32,
         )
 
         image = self.decode_latents(latents, output_type=output_type)
+        self.maybe_free_model_hooks()
         if not return_dict:
             return (image, latents)
         return ImagePipelineOutput(images=image)

edm2-img512-s-fid/pipeline.py

@@ -350,6 +350,9 @@ class EDM2Pipeline(DiffusionPipeline):
 
         device = self._execution_device
         dtype = self.unet.dtype
+        self.unet.eval()
+        if getattr(self, "gnet", None) is not None:
+            self.gnet.eval()
         labels = self._normalize_class_labels(class_labels, batch_size=batch_size, device=device)
         noise = self.prepare_latents(batch_size, height, width, dtype, device, generator)
 
@@ -375,13 +378,14 @@ class EDM2Pipeline(DiffusionPipeline):
         latents = self._sample_edm2_heun(
             denoise_fn=denoise_fn,
             noise=noise,
-            sigmas=self.scheduler.sigmas.to(device),
+            sigmas=self.scheduler.sigmas.to(device).clone(),
             generator=generator,
             progress_bar=self.progress_bar,
             dtype=torch.float32,
         )
 
         image = self.decode_latents(latents, output_type=output_type)
+        self.maybe_free_model_hooks()
         if not return_dict:
             return (image, latents)
         return ImagePipelineOutput(images=image)

edm2-img512-xl-fid/pipeline.py

@@ -350,6 +350,9 @@ class EDM2Pipeline(DiffusionPipeline):
 
         device = self._execution_device
         dtype = self.unet.dtype
+        self.unet.eval()
+        if getattr(self, "gnet", None) is not None:
+            self.gnet.eval()
         labels = self._normalize_class_labels(class_labels, batch_size=batch_size, device=device)
         noise = self.prepare_latents(batch_size, height, width, dtype, device, generator)
 
@@ -375,13 +378,14 @@ class EDM2Pipeline(DiffusionPipeline):
         latents = self._sample_edm2_heun(
             denoise_fn=denoise_fn,
             noise=noise,
-            sigmas=self.scheduler.sigmas.to(device),
+            sigmas=self.scheduler.sigmas.to(device).clone(),
             generator=generator,
             progress_bar=self.progress_bar,
             dtype=torch.float32,
         )
 
         image = self.decode_latents(latents, output_type=output_type)
+        self.maybe_free_model_hooks()
         if not return_dict:
             return (image, latents)
         return ImagePipelineOutput(images=image)

edm2-img512-xs-fid/pipeline.py

@@ -350,6 +350,9 @@ class EDM2Pipeline(DiffusionPipeline):
 
         device = self._execution_device
         dtype = self.unet.dtype
+        self.unet.eval()
+        if getattr(self, "gnet", None) is not None:
+            self.gnet.eval()
         labels = self._normalize_class_labels(class_labels, batch_size=batch_size, device=device)
         noise = self.prepare_latents(batch_size, height, width, dtype, device, generator)
 
@@ -375,13 +378,14 @@ class EDM2Pipeline(DiffusionPipeline):
         latents = self._sample_edm2_heun(
             denoise_fn=denoise_fn,
             noise=noise,
-            sigmas=self.scheduler.sigmas.to(device),
+            sigmas=self.scheduler.sigmas.to(device).clone(),
             generator=generator,
             progress_bar=self.progress_bar,
             dtype=torch.float32,
         )
 
         image = self.decode_latents(latents, output_type=output_type)
+        self.maybe_free_model_hooks()
         if not return_dict:
             return (image, latents)
         return ImagePipelineOutput(images=image)

edm2-img512-xxl-fid/pipeline.py

@@ -350,6 +350,9 @@ class EDM2Pipeline(DiffusionPipeline):
 
         device = self._execution_device
         dtype = self.unet.dtype
+        self.unet.eval()
+        if getattr(self, "gnet", None) is not None:
+            self.gnet.eval()
         labels = self._normalize_class_labels(class_labels, batch_size=batch_size, device=device)
         noise = self.prepare_latents(batch_size, height, width, dtype, device, generator)
 
@@ -375,13 +378,14 @@ class EDM2Pipeline(DiffusionPipeline):
         latents = self._sample_edm2_heun(
             denoise_fn=denoise_fn,
             noise=noise,
-            sigmas=self.scheduler.sigmas.to(device),
+            sigmas=self.scheduler.sigmas.to(device).clone(),
             generator=generator,
             progress_bar=self.progress_bar,
             dtype=torch.float32,
         )
 
         image = self.decode_latents(latents, output_type=output_type)
+        self.maybe_free_model_hooks()
         if not return_dict:
             return (image, latents)
         return ImagePipelineOutput(images=image)