add lic & readme, code reformat & review

Author: Kahsolt

Euler.py

@@ -1,22 +1,28 @@
-from diffusers import EulerDiscreteScheduler
-from torch import Tensor
-import torch
 from typing import Callable, List, Optional, Tuple, Union, Dict, Any, Literal
+
+import torch
+from torch import Tensor
+from diffusers import EulerDiscreteScheduler
 from diffusers.utils import randn_tensor
 from diffusers.configuration_utils import ConfigMixin
 from diffusers.schedulers.scheduling_utils import SchedulerMixin
+
+
 class Euler(EulerDiscreteScheduler, SchedulerMixin, ConfigMixin):
-    history_d=0
-    momentum=0.95
-    momentum_hist=0.75
-    def init_hist_d(self,x:Tensor) -> Union[Literal[0], Tensor]:
+
+    history_d = 0
+    momentum = 0.95
+    momentum_hist = 0.75
+
+    def init_hist_d(self, x:Tensor) -> Union[Literal[0], Tensor]:
         # memorize delta momentum
-        if   self.history_d == 0:      self.history_d = 0
+        if   self.history_d == 0:           self.history_d = 0
         elif self.history_d == 'rand_init': self.history_d = x
         elif self.history_d == 'rand_new':  self.history_d = torch.randn_like(x)
         else: raise ValueError(f'unknown momentum_hist_init: {self.history_d}')
+
     def momentum_step(self, x:Tensor, d:Tensor, dt:Tensor):
-        hd=self.history_d
+        hd = self.history_d
         # correct current `d` with momentum
         p = 1.0 - self.momentum
         self.momentum_d = (1.0 - p) * d + p * hd    
@@ -30,8 +36,10 @@ def momentum_step(self, x:Tensor, d:Tensor, dt:Tensor):
             hd = self.momentum_d
         else:
             hd = (1.0 - q) * hd + q * self.momentum_d
-        self.history_d=hd
+        self.history_d = hd
+
         return x
+
     def step(
         self,
         model_output: torch.FloatTensor,
@@ -124,10 +132,8 @@ def step(
         derivative = (sample - pred_original_sample) / sigma_hat
 
         dt = self.sigmas[step_index + 1] - sigma_hat
-
-        prev_sample = self.momentum_step(sample,derivative,dt)
-        if not return_dict:
-            return (prev_sample,)
-
-        output={prev_sample:prev_sample, pred_original_sample:pred_original_sample}
-        return output
+        prev_sample = self.momentum_step(sample, derivative, dt)
+        
+        if not return_dict: return (prev_sample,)
+        output = { prev_sample: prev_sample, pred_original_sample: pred_original_sample }
+        return output

EulerA.py

@@ -1,22 +1,28 @@
-from diffusers import EulerAncestralDiscreteScheduler
-from torch import Tensor
-import torch
 from typing import Callable, List, Optional, Tuple, Union, Dict, Any, Literal
+
+import torch
+from torch import Tensor
+from diffusers import EulerAncestralDiscreteScheduler
 from diffusers.utils import randn_tensor
 from diffusers.configuration_utils import ConfigMixin
 from diffusers.schedulers.scheduling_utils import SchedulerMixin
+
+
 class EulerA(EulerAncestralDiscreteScheduler, SchedulerMixin, ConfigMixin):
-    history_d=0
-    momentum=0.95
-    momentum_hist=0.75
-    def init_hist_d(self,x:Tensor) -> Union[Literal[0], Tensor]:
+
+    history_d = 0
+    momentum = 0.95
+    momentum_hist = 0.75
+
+    def init_hist_d(self, x:Tensor) -> Union[Literal[0], Tensor]:
         # memorize delta momentum
-        if   self.history_d == 0:      self.history_d = 0
+        if   self.history_d == 0:           self.history_d = 0
         elif self.history_d == 'rand_init': self.history_d = x
         elif self.history_d == 'rand_new':  self.history_d = torch.randn_like(x)
         else: raise ValueError(f'unknown momentum_hist_init: {self.history_d}')
+    
     def momentum_step(self, x:Tensor, d:Tensor, dt:Tensor):
-        hd=self.history_d
+        hd = self.history_d
         # correct current `d` with momentum
         p = 1.0 - self.momentum
         self.momentum_d = (1.0 - p) * d + p * hd    
@@ -30,8 +36,10 @@ def momentum_step(self, x:Tensor, d:Tensor, dt:Tensor):
             hd = self.momentum_d
         else:
             hd = (1.0 - q) * hd + q * self.momentum_d
-        self.history_d=hd
+        self.history_d = hd
+
         return x
+    
     def step(
         self,
         model_output: torch.FloatTensor,
@@ -101,15 +109,13 @@ def step(
         derivative = (sample - pred_original_sample) / sigma
 
         dt = sigma_down - sigma
-        
-        prev_sample = self.momentum_step(sample,derivative,dt)
+        prev_sample = self.momentum_step(sample, derivative, dt)
 
         device = model_output.device
         noise = randn_tensor(model_output.shape, dtype=model_output.dtype, device=device, generator=generator)
 
         prev_sample = prev_sample + noise * sigma_up
 
-        if not return_dict:
-            return (prev_sample,)
-        output={prev_sample:prev_sample, pred_original_sample:pred_original_sample}
-        return output
+        if not return_dict: return (prev_sample,)
+        output = { prev_sample: prev_sample, pred_original_sample: pred_original_sample }
+        return output

LICENSE

@@ -0,0 +1,19 @@
+MIT License
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -0,0 +1,5 @@
+# modified-euler-samplers-for-sonar-diffusers
+
+----
+
+This repo ports the **momentum mechanism** of [sd-webui-sonar](https://github.com/Kahsolt/stable-diffusion-webui-sonar) on `Euler` and `Euler a` sampler to [huggingface/diffusers](https://github.com/huggingface/diffusers). 

example (partial).py

@@ -0,0 +1,48 @@
+from schedulers.EulerA import EulerA
+# Initialize the Celery app
+
+controlnet = ControlNetModel.from_pretrained(
+  "lllyasviel/control_v11p_sd15_openpose", 
+  torch_dtype=torch.float16,
+  local_files_only=True,
+)
+pipe = StableDiffusionControlNetPipeline.from_pretrained(
+  "runwayml/stable-diffusion-v1-5",
+  controlnet=controlnet,
+  local_files_only=True, 
+  torch_dtype=torch.float16,
+  safety_checker=None, 
+  requires_safety_checker=False,
+).to('cuda')
+
+# import one of the 2 schedulers from this repo
+pipe.scheduler = EulerA.from_config(pipe.scheduler.config)
+
+# choose from [0, 'rand_new', 'rand_init']
+pipe.scheduler.history_d = 'rand_new'
+# number should be between -1 and 1
+pipe.scheduler.momentum = 0.95
+# number should be between -1 and 1
+pipe.scheduler.momentum_hist = 0.75
+
+buffer = open('img0.png', 'rb')
+buffer.seek(0)
+image_bytes = buffer.read()
+images = Image.open(BytesIO(image_bytes))
+
+start_time = time.time()
+generator = torch.manual_seed(2733424006)
+image=pipe(
+  "A person standing in a field of flowers, 4k, realistic",
+  images,
+  num_inference_steps=20,
+  height=512,
+  width=512,
+  generator=generator
+).images[0]
+end_time = time.time()
+execution_time = end_time - start_time
+print("Execution time: {:.2f} seconds".format(execution_time))
+
+# print(image)
+image.save('img1.png', format='PNG')