general normal distribution

README.md

@@ -14,7 +14,10 @@ Currently implemented methods:
   Example: `custom_pipeline.load_initial_noise_modifier(method="fixed-seed", seed=…)`
 -  Golden Noise for Diffusion Models: A Learning Framework (Zhou et al., https://arxiv.org/abs/2411.09502). 
   Example: `custom_pipeline.load_initial_noise_modifier(method="golden-noise", npnet_path=…)`
-Demo Notebook: [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1-owYN8r2TbT-Je_eTEpnIMLj1nvxPYqI?usp=sharing)
+
+
+
+Demo Notebook: [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1-owYN8r2TbT-Je_eTEpnIMLj1nvxPYqI#scrollTo=HQS6OQ44jz66)
 
 ## Citation 
 

pipeline.py

@@ -85,6 +85,25 @@ class NoiseLoaderMixin:
         elif method == "fixed-seed":
             # This is for demonstration purposes to see the impact of different seeds
             self.seed = method_args.get("seed", 0)
+        elif method == "general-normal-distribution":
+            self.init_noise_mean = method_args.get("init_noise_mean", [0.0, 0.0, 0.0, 0.0])
+            self.init_noise_std = method_args.get("init_noise_std", [1.0, 1.0, 1.0, 1.0])
+
+            if isinstance(self.init_noise_mean, (list, tuple)):
+                if len(self.init_noise_mean) != 4:
+                    raise ValueError("Mean must be a scalar or a list/tuple of 4 values (one per channel)")
+                self.init_noise_mean = torch.tensor(self.init_noise_mean, dtype=self.dtype, device=self.device).view(1, 4, 1, 1)
+            else:
+                # Scalar case - broadcast to all channels
+                self.init_noise_mean = torch.tensor([self.init_noise_mean] * 4, dtype=self.dtype, device=self.device).view(1, 4, 1, 1)
+
+            if isinstance(self.init_noise_std, (list, tuple)):
+                if len(self.init_noise_std) != 4:
+                    raise ValueError("Std must be a scalar or a list/tuple of 4 values (one per channel)")
+                self.init_noise_std = torch.tensor(self.init_noise_std, dtype=self.dtype, device=self.device).view(1, 4, 1, 1)
+            else:
+                # Scalar case - broadcast to all channels
+                self.init_noise_std = torch.tensor([self.init_noise_std] * 4, dtype=self.dtype, device=self.device).view(1, 4, 1, 1)
         else:
             raise NotImplementedError(f"No initial noise method is implemented for the given method {method}")
 
@@ -117,6 +136,12 @@ class NoiseLoaderMixin:
         elif self.initial_noise_sampling_method == "fixed-seed":
             generator = torch.Generator(device=self.device).manual_seed(self.seed)
             return torch.randn(shape, generator=generator, dtype=self.dtype, device=self.device)
+        elif self.initial_noise_sampling_method == "general-normal-distribution":
+            # Generate signal using custom mean and std
+            initial_noise = self.init_noise_mean + self.init_noise_std * randn_tensor(
+                shape, generator=generator, device=self.device, dtype=self.dtype
+            )
+            return initial_noise
         
     def unload_initial_noise_modifier(self):
         """