Update sd_simple_kes_v3_fix?/simple_kes_v3.py

sd_simple_kes_v3_fix?/simple_kes_v3.py

@@ -922,7 +922,14 @@ class SimpleKEScheduler:
             self.change = torch.abs(self.sigs[i] - self.sigs[i - 1])
             # Safely extract scalar for both tensor and float
             self.change_log.append(self.extract_scalar(self.change))
-            relative_sigma_progress = (self.blended_sigma - self.sigs[-1].item()) / self.blended_sigma                            
+            # --- ИСПРАВЛЕНИЕ ДЕЛЕНИЯ НА НОЛЬ (v2) ---
+            blended_scalar = self.extract_scalar(self.blended_sigma)
+            if blended_scalar != 0:
+                relative_sigma_progress = (blended_scalar - self.sigs[-1].item()) / blended_scalar
+            else:
+                relative_sigma_progress = 0.0
+            # --- КОНЕЦ ИСПРАВЛЕНИЯ ---
+            
             recent_changes = torch.abs(torch.tensor(self.change_log[-5:]))
             max_change = torch.max(recent_changes).item()                        
             mean_change = torch.mean(recent_changes).item()
@@ -963,13 +970,20 @@ class SimpleKEScheduler:
                 # Start checking for early stopping after minimum steps
                 if i > self.safety_minimum_stop_step and len(self.change_log) > 10:                  
                     # Calculate variance and dynamic threshold
-                    self.blended_tensor = torch.tensor(self.prepass_blended_sigmas) 
-                    if self.device == 'cpu':
-                        self.sigma_variance = np.var(self.prepass_blended_sigmas)
+                    self.blended_tensor = torch.tensor(self.prepass_blended_sigmas, device=self.device) # <-- Добавлен device
+                    
+                    # --- ИСПРАВЛЕНИЕ РАСЧЕТА ДИСПЕРСИИ ---
+                    # self.sigs еще не заполнен (полон нулей). 
+                    # Мы должны использовать 'blended_tensor', который содержит реальные значения.
+                    if self.device.type == 'cpu':
+                        # np.var работает на списке Python
+                        self.sigma_variance = np.var(self.prepass_blended_sigmas) 
                     else: 
-                        self.sigma_variance = torch.var(self.sigs).item()
+                        # torch.var работает на тензоре
+                        self.sigma_variance = torch.var(self.blended_tensor).item() 
+                    # --- КОНЕЦ ИСПРАВЛЕНИЯ ---
 
-                    self.min_sigma_threshold = self.sigma_variance * self.sigma_variance_scale  # scale factor can be tuned
+                    self.min_sigma_threshold = self.sigma_variance * self.sigma_variance_scale
                     self.prepass_log(f"\n--- Early Stopping Evaluation at Step {i} ---")
                     self.prepass_log(f"Current Blended Prepass Sigma: {self.prepass_blended_sigma:.6f}")
                     self.prepass_log(f"Sigma Variance: {self.sigma_variance:.6f}")
@@ -1061,10 +1075,11 @@ class SimpleKEScheduler:
                 # Early Stopping Evaluation
                 if i > self.safety_minimum_stop_step and len(self.change_log) > 5:                                 
                     final_target_sigma = self.sigs[-1].item()  # or use min(self.sigmas) if preferred                      
-                    if self.blended_sigma != 0:
-                        relative_sigma_progress = (self.blended_sigma - final_target_sigma) / self.blended_sigma
+                    blended_scalar = self.extract_scalar(self.blended_sigma)
+                    if blended_scalar != 0:
+                        relative_sigma_progress = (blended_scalar - self.sigs[-1].item()) / blended_scalar
                     else:
-                        relative_sigma_progress = 0  # Assume fully converged if blended_sigma is 0                                    
+                        relative_sigma_progress = 0.0                                   
                     # Optional: Show variance but no need to stop on it
                     self.sigma_variance = torch.var(self.sigs).item() if self.device != 'cpu' else np.var(self.blended_sigmas)
                     self.log(f"Sigma Variance: {self.sigma_variance:.6f}")
@@ -1920,36 +1935,46 @@ class SimpleKEScheduler:
             blend_weights = self.blend_weights
             
         )
-        self.sigma_variance = torch.var(self.sigs).item()          
-        if self.sharpen_mode in ['last_n', 'both']:
+        self.sigma_variance = torch.var(self.sigs).item()
+        # --- НАЧАЛО ИСПРАВЛЕНИЯ ДВОЙНОЙ РЕЗКОСТИ ---
+        # Логика для 'full': применяется только 'full'
+        if self.sharpen_mode == 'full':
+            self.sharpen_mask = torch.where(self.sigs < self.sigma_min * 1.5, self.sharpness, 1.0).to(self.device)
+            sharpen_indices = torch.where(self.sharpen_mask < 1.0)[0].tolist()
+            self.sigs = self.sigs * self.sharpen_mask
+            self.log(f"[Sharpen Mask] 'full' sharpening applied at steps: {sharpen_indices}")
+
+        # Логика для 'last_n': применяется только 'last_n'
+        elif self.sharpen_mode == 'last_n':
+            # Примечание: эта проверка дисперсии взята из оригинальной (странной) логики
             if self.sigma_variance < self.sharpen_variance_threshold:
-                # Apply full sharpening
-                self.sharpen_mask = torch.where(self.sigs < self.sigma_min * 1.5, self.sharpness, 1.0).to(self.device)
-                sharpen_indices = torch.where(self.sharpen_mask < 1.0)[0].tolist()
-                self.sigs = self.sigs * self.sharpen_mask
-                self.log(f"[Sharpen Mask] Full sharpening applied (low variance). Steps: {sharpen_indices}")
+                 self.log(f"[Sharpen Mask] 'last_n' mode, but variance is low ({self.sigma_variance:.6f} < {self.sharpen_variance_threshold}). Skipping 'last_n'.")
             else:
-                # Apply sharpening only to the last N steps
                 recent_sigs = self.sigs[-self.sharpen_last_n_steps:]
                 sharpen_mask = torch.where(recent_sigs < self.sigma_min * 1.5, self.sharpness, 1.0).to(self.device)
                 sharpen_indices = torch.where(sharpen_mask < 1.0)[0].tolist()
                 self.sigs[-self.sharpen_last_n_steps:] = recent_sigs * sharpen_mask
+                self.log(f"[Sharpen Mask] 'last_n' sharpening applied at steps: {sharpen_indices}")
 
-                # Now loop per step if desired (safely inside this block)
-                for j in range(len(self.sigs) - self.sharpen_last_n_steps, len(self.sigs)):
-                    if self.sigs[j] < self.sigma_min * 1.5:
-                        old_value = self.sigs[j].item()
-                        self.sigs[j] = self.sigs[j] * self.sharpness
-                        self.log(f"[Sharpening] Step {j+1}: Applied sharpening. Sigma changed from {old_value:.6f} to {self.sigs[j].item():.6f}")
-                    else:
-                        self.log(f"[Sharpening] Step {j+1}: No sharpening applied. Sigma: {self.sigs[j].item():.6f}")
-
-        if self.sharpen_mode in ['full', 'both']:
-            # Optional: Additional full sharpening (if needed)
+        # Логика для 'both': применяется 'full', А ЗАТЕМ 'last_n'
+        elif self.sharpen_mode == 'both':
+            # 1. Применяем 'full'
             self.sharpen_mask = torch.where(self.sigs < self.sigma_min * 1.5, self.sharpness, 1.0).to(self.device)
-            sharpen_indices = torch.where(self.sharpen_mask < 1.0)[0].tolist()
+            sharpen_indices_full = torch.where(self.sharpen_mask < 1.0)[0].tolist()
             self.sigs = self.sigs * self.sharpen_mask
-            self.log(f"[Sharpen Mask] Full sharpening applied at steps: {sharpen_indices}")
+            self.log(f"[Sharpen Mask] 'both' (part 1/2): 'full' sharpening applied at steps: {sharpen_indices_full}")
+
+            # 2. Применяем 'last_n' (с той же проверкой дисперсии)
+            if self.sigma_variance < self.sharpen_variance_threshold:
+                 self.log(f"[Sharpen Mask] 'both' (part 2/2): variance is low ({self.sigma_variance:.6f} < {self.sharpen_variance_threshold}). Skipping 'last_n' part.")
+            else:
+                recent_sigs = self.sigs[-self.sharpen_last_n_steps:]
+                sharpen_mask = torch.where(recent_sigs < self.sigma_min * 1.5, self.sharpness, 1.0).to(self.device)
+                sharpen_indices_last_n = torch.where(sharpen_mask < 1.0)[0].tolist()
+                self.sigs[-self.sharpen_last_n_steps:] = recent_sigs * sharpen_mask
+                self.log(f"[Sharpen Mask] 'both' (part 2/2): 'last_n' sharpening applied at steps: {sharpen_indices_last_n}")
+        
+        # --- КОНЕЦ ИСПРАВЛЕНИЯ ДВОЙНОЙ РЕЗКОСТИ ---
         
         '''
         sigma_hash = self.generate_sigma_hash(steps, sigma_min, sigma_max, rho, device, schedule_type, decay_pattern, suffix=None)