Update pipelines/pipeline_seesr.py

pipelines/pipeline_seesr.py

@@ -778,7 +778,6 @@ class StableDiffusionControlNetPipeline(DiffusionPipeline, TextualInversionLoade
         return torch.tile(torch.tensor(weights, device=self.device), (nbatches, self.unet.config.in_channels, 1, 1))
 
     @perfcount
-    @torch.no_grad()
     @replace_example_docstring(EXAMPLE_DOC_STRING)
     def __call__(
         self,
@@ -808,7 +807,8 @@ class StableDiffusionControlNetPipeline(DiffusionPipeline, TextualInversionLoade
         ram_encoder_hidden_states=None,
         latent_tiled_size=320,
         latent_tiled_overlap=4,
-        args=None
+        use_KDS=True,
+        args=None,
     ):
         r"""
         Function invoked when calling the pipeline for generation.
@@ -996,6 +996,9 @@ class StableDiffusionControlNetPipeline(DiffusionPipeline, TextualInversionLoade
         # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
         extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
 
+        if use_KDS:
+            latents.requires_grad_(True)
+
         # 8. Denoising loop
         num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
         with self.progress_bar(total=num_inference_steps) as progress_bar:
@@ -1183,9 +1186,39 @@ class StableDiffusionControlNetPipeline(DiffusionPipeline, TextualInversionLoade
                     noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
 
 
+                if use_KDS:
 
-                # compute the previous noisy sample x_t -> x_t-1
-                latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
+                    # 2) Compute x₀ prediction
+                    beta_t   = 1 - self.scheduler.alphas_cumprod[t]
+                    alpha_t  = self.scheduler.alphas_cumprod[t].sqrt()
+                    sigma_t  = beta_t.sqrt()
+                    x0_pred  = (latents - sigma_t * noise_pred) / alpha_t
+                    
+                    # 3) Apply KDE steering
+                    m_shift   = kde_grad(x0_pred)
+                    delta_t   = gamma_0 * (1 - i / (len(timesteps_tensor) - 1))
+                    x0_steer  = x0_pred + delta_t * m_shift
+                    
+                    # 4) Recompute “noise” for DDIM step
+                    noise_pred_kds = (latents - alpha_t * x0_steer) / sigma_t
+                    
+                    # 5) Determine prev alphas
+                    if i < len(timesteps_tensor) - 1:
+                      next_t = timesteps_tensor[i + 1]
+                      alpha_prev = self.scheduler.alphas_cumprod[next_t].sqrt()
+                    else:
+                      alpha_prev = self.scheduler.final_alpha_cumprod.sqrt()
+                    
+                    sigma_prev = (1 - alpha_prev**2).sqrt()
+                    
+                    # 6) Form next latent per DDIM
+                    latents = (
+                      alpha_prev * x0_steer
+                      + sigma_prev * noise_pred_kds
+                    ).detach().requires_grad_(True)
+                else:
+                    # compute the previous noisy sample x_t -> x_t-1
+                    latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0].requires_grad_(False)
 
                 # call the callback, if provided
                 if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):