improve: less ram usage

Author: samedii

posthoc_ema/karras_ema.py

@@ -38,7 +38,12 @@ def inplace_lerp(tgt: Tensor, src: Tensor, weight):
         src: Source tensor to interpolate towards
         weight: Interpolation weight between 0 and 1
     """
-    tgt.lerp_(src.to(tgt.device), weight)
+    # Check if tensor is integer type - integer tensors can't use lerp
+    # but we want to silently handle them instead of raising errors
+    if tgt.dtype in [torch.int, torch.int8, torch.int16, torch.int32, torch.int64, torch.long]:
+        tgt.copy_(src.to(tgt.device))
+    else:
+        tgt.lerp_(src.to(tgt.device), weight)
 
 
 class KarrasEMA(Module):
@@ -49,94 +54,73 @@ class KarrasEMA(Module):
         model: Model to create EMA of
         sigma_rel: Relative standard deviation for EMA profile
         gamma: Alternative parameterization via gamma (don't specify both)
-        ema_model: Optional pre-initialized EMA model
         update_every: Number of steps between EMA updates
         frozen: Whether to freeze EMA updates
         param_or_buffer_names_no_ema: Parameter/buffer names to exclude from EMA
         ignore_names: Parameter/buffer names to ignore
         ignore_startswith_names: Parameter/buffer name prefixes to ignore
         only_save_diff: If True, only save parameters with requires_grad=True
+        device: Device to store EMA parameters on (default='cpu')
     """
 
+    # Buffers that should always be included in the state dict even with only_save_diff=True
+    _ALWAYS_INCLUDE_BUFFERS = {"running_mean", "running_var", "num_batches_tracked"}
+    
     def __init__(
         self,
         model: Module,
         sigma_rel: float | None = None,
         gamma: float | None = None,
-        ema_model: Module | Callable[[], Module] | None = None,
         update_every: int = 10,
         frozen: bool = False,
         param_or_buffer_names_no_ema: set[str] = set(),
         ignore_names: set[str] = set(),
         ignore_startswith_names: set[str] = set(),
         only_save_diff: bool = False,
+        device: str = 'cpu',
     ):
         super().__init__()
-
-        assert exists(sigma_rel) ^ exists(
-            gamma
-        ), "either sigma_rel or gamma must be given"
-
-        if exists(sigma_rel):
-            gamma = sigma_rel_to_gamma(sigma_rel)
-
+        
+        # Store all the configuration parameters first
         self.gamma = gamma
         self.frozen = frozen
         self.update_every = update_every
         self.only_save_diff = only_save_diff
-
+        self.ignore_names = ignore_names
+        self.ignore_startswith_names = ignore_startswith_names
+        self.param_or_buffer_names_no_ema = param_or_buffer_names_no_ema
+        self.device = device
+        
+        assert exists(sigma_rel) ^ exists(gamma), "either sigma_rel or gamma must be given"
+        
+        if exists(sigma_rel):
+            gamma = sigma_rel_to_gamma(sigma_rel)
+            self.gamma = gamma
+        
         # Store reference to online model
         self.online_model = [model]
-
-        # Initialize EMA model
-        if callable(ema_model) and not isinstance(ema_model, Module):
-            ema_model = ema_model()
-
-        # Store original device
-        original_device = next(model.parameters()).device
-
-        # Move model to CPU before copying to avoid VRAM spike
-        model.cpu()
-
-        try:
-            # Create EMA model on CPU
-            self.ema_model = (ema_model if exists(ema_model) else deepcopy(model)).cpu()
-
-            # Ensure all parameters and buffers are on CPU and detached
-            for p in self.ema_model.parameters():
-                p.data = p.data.cpu().detach()
-            for b in self.ema_model.buffers():
-                b.data = b.data.cpu().detach()
-
-            # Move model back to original device
-            model.to(original_device)
-
-            # Get parameter names for floating point or complex parameters
-            self.param_names = {
-                name
-                for name, param in self.ema_model.named_parameters()
-                if torch.is_floating_point(param) or torch.is_complex(param)
-            }
-
-            # Get buffer names for floating point or complex buffers
-            self.buffer_names = {
-                name
-                for name, buffer in self.ema_model.named_buffers()
-                if torch.is_floating_point(buffer) or torch.is_complex(buffer)
-            }
-
-            # Names to ignore
-            self.param_or_buffer_names_no_ema = param_or_buffer_names_no_ema
-            self.ignore_names = ignore_names
-            self.ignore_startswith_names = ignore_startswith_names
-
-            # State buffers on CPU
-            self.register_buffer("initted", torch.tensor(False, device="cpu"))
-            self.register_buffer("step", torch.tensor(0, device="cpu"))
-        except:
-            # Ensure model is moved back even if initialization fails
-            model.to(original_device)
-            raise
+        
+        # Instead of copying the whole model, just store parameter tensors
+        self.ema_params = {}
+        self.ema_buffers = {}
+        
+        # Get parameter and buffer names to track
+        with torch.no_grad():
+            for name, param in model.named_parameters():
+                if self._should_update_param(name):
+                    if not only_save_diff or param.requires_grad:
+                        self.ema_params[name] = param.detach().clone().to(self.device)
+                        
+            for name, buffer in model.named_buffers():
+                if self._should_update_param(name):
+                    buffer_name = name.split('.')[-1]  # Get the base name
+                    # Always include critical buffers regardless of only_save_diff
+                    if not only_save_diff or buffer.requires_grad or buffer_name in self._ALWAYS_INCLUDE_BUFFERS:
+                        self.ema_buffers[name] = buffer.detach().clone().to(self.device)
+        
+        # State buffers
+        self.register_buffer("initted", torch.tensor(False))
+        self.register_buffer("step", torch.tensor(0))
 
     @property
     def beta(self):
@@ -161,42 +145,33 @@ def update(self):
     def copy_params_from_model_to_ema(self):
         """Copy parameters from online model to EMA model."""
         # Copy parameters
-        for (name, ma_params), (_, current_params) in zip(
-            self.get_params_iter(self.ema_model),
-            self.get_params_iter(self.online_model[0]),
-        ):
-            if self._should_update_param(name):
-                inplace_copy(ma_params.data, current_params.data)
-
-        # Copy buffers
-        for (name, ma_buffer), (_, current_buffer) in zip(
-            self.get_buffers_iter(self.ema_model),
-            self.get_buffers_iter(self.online_model[0]),
-        ):
-            if self._should_update_param(name):
-                inplace_copy(ma_buffer.data, current_buffer.data)
+        with torch.no_grad():
+            for name, param in self.online_model[0].named_parameters():
+                if name in self.ema_params:
+                    # Explicitly move to device (usually CPU)
+                    self.ema_params[name] = param.detach().clone().to(self.device)
+
+            # Copy buffers
+            for name, buffer in self.online_model[0].named_buffers():
+                if name in self.ema_buffers:
+                    # Explicitly move to device (usually CPU)
+                    self.ema_buffers[name] = buffer.detach().clone().to(self.device)
 
     def update_moving_average(self):
         """Update EMA weights using current beta value."""
         current_decay = self.beta
 
-        # Update parameters
-        for (name, current_params), (_, ma_params) in zip(
-            self.get_params_iter(self.online_model[0]),
-            self.get_params_iter(self.ema_model),
-        ):
-            if not self._should_update_param(name):
-                continue
-            inplace_lerp(ma_params.data, current_params.data, 1.0 - current_decay)
-
-        # Update buffers
-        for (name, current_buffer), (_, ma_buffer) in zip(
-            self.get_buffers_iter(self.online_model[0]),
-            self.get_buffers_iter(self.ema_model),
-        ):
-            if not self._should_update_param(name):
-                continue
-            inplace_lerp(ma_buffer.data, current_buffer.data, 1.0 - current_decay)
+        # Update parameters using the simplified lerp function (which now handles integer tensors)
+        for name, current_params in self.online_model[0].named_parameters():
+            if name in self.ema_params:
+                # inplace_lerp now handles integer tensors internally
+                inplace_lerp(self.ema_params[name], current_params.data, 1.0 - current_decay)
+
+        # Update buffers with the same simplified approach
+        for name, current_buffer in self.online_model[0].named_buffers():
+            if name in self.ema_buffers:
+                # inplace_lerp now handles integer tensors internally
+                inplace_lerp(self.ema_buffers[name], current_buffer.data, 1.0 - current_decay)
 
     def _should_update_param(self, name: str) -> bool:
         """Check if parameter should be updated based on ignore rules."""
@@ -208,10 +183,17 @@ def _should_update_param(self, name: str) -> bool:
             return False
         return True
 
+    def _parameter_requires_grad(self, name: str) -> bool:
+        """Check if parameter requires gradients in the online model."""
+        for n, p in self.online_model[0].named_parameters():
+            if n == name:
+                return p.requires_grad
+        return False
+
     def get_params_iter(self, model):
         """Get iterator over model's parameters."""
         for name, param in model.named_parameters():
-            if name not in self.param_names:
+            if name not in self.ema_params:
                 continue
             if self.only_save_diff and not param.requires_grad:
                 continue
@@ -220,17 +202,19 @@ def get_params_iter(self, model):
     def get_buffers_iter(self, model):
         """Get iterator over model's buffers."""
         for name, buffer in model.named_buffers():
-            if name not in self.buffer_names:
+            if name not in self.ema_buffers:
                 continue
-            if self.only_save_diff and not buffer.requires_grad:
+            
+            # Handle critical buffers that should always be included
+            buffer_name = name.split('.')[-1]
+            if self.only_save_diff and not buffer.requires_grad and buffer_name not in self._ALWAYS_INCLUDE_BUFFERS:
                 continue
+            
             yield name, buffer
 
     def iter_all_ema_params_and_buffers(self):
         """Get iterator over all EMA parameters and buffers."""
-        for name, param in self.ema_model.named_parameters():
-            if name not in self.param_names:
-                continue
+        for name, param in self.ema_params.items():
             if name in self.param_or_buffer_names_no_ema:
                 continue
             if name in self.ignore_names:
@@ -239,21 +223,10 @@ def iter_all_ema_params_and_buffers(self):
                 continue
             yield param
 
-        for name, buffer in self.ema_model.named_buffers():
-            if name not in self.buffer_names:
-                continue
-            if name in self.param_or_buffer_names_no_ema:
-                continue
-            if name in self.ignore_names:
-                continue
-            if any(name.startswith(prefix) for prefix in self.ignore_startswith_names):
-                continue
-            yield buffer
-
     def iter_all_model_params_and_buffers(self, model: Module):
         """Get iterator over all model parameters and buffers."""
         for name, param in model.named_parameters():
-            if name not in self.param_names:
+            if name not in self.ema_params:
                 continue
             if name in self.param_or_buffer_names_no_ema:
                 continue
@@ -263,59 +236,66 @@ def iter_all_model_params_and_buffers(self, model: Module):
                 continue
             yield param
 
-        for name, buffer in model.named_buffers():
-            if name not in self.buffer_names:
-                continue
-            if name in self.param_or_buffer_names_no_ema:
-                continue
-            if name in self.ignore_names:
-                continue
-            if any(name.startswith(prefix) for prefix in self.ignore_startswith_names):
-                continue
-            yield buffer
-
     def __call__(self, *args, **kwargs):
         """Forward pass using EMA model."""
-        return self.ema_model(*args, **kwargs)
+        raise NotImplementedError("KarrasEMA no longer maintains a full model copy")
+
+    @property
+    def ema_model(self):
+        """
+        For backward compatibility with tests.
+        Creates a temporary model with EMA parameters.
+        
+        Returns:
+            Module: A copy of the online model with EMA parameters
+        """
+        # Create a copy of the online model
+        model_copy = deepcopy(self.online_model[0])
+        
+        # Load EMA parameters into the model
+        for name, param in model_copy.named_parameters():
+            if name in self.ema_params:
+                param.data.copy_(self.ema_params[name])
+            
+        # Load EMA buffers into the model
+        for name, buffer in model_copy.named_buffers():
+            if name in self.ema_buffers:
+                buffer.data.copy_(self.ema_buffers[name])
+        
+        # Ensure the model is on CPU
+        model_copy.to('cpu')
+        return model_copy
 
     def state_dict(self):
-        """Get state dict for EMA model."""
+        """Get state dict with EMA parameters."""
         state_dict = {}
-
-        # Save parameters based on only_save_diff flag
-        for name, param in self.ema_model.named_parameters():
-            if name not in self.param_names:
-                continue
-            if self.only_save_diff and not param.requires_grad:
-                continue
-            state_dict[name] = param.data
-
-        # Save buffers
-        for name, buffer in self.ema_model.named_buffers():
-            if name not in self.buffer_names:
-                continue
-            state_dict[name] = buffer.data
-
-        # Save internal state
-        state_dict["initted"] = self.initted.data
-        state_dict["step"] = self.step.data
-
+        
+        # For parameters, respect only_save_diff
+        for name, param in self.ema_params.items():
+            if not self.only_save_diff or self._parameter_requires_grad(name):
+                state_dict[name] = param.data
+        
+        # For buffers, identify which ones should always be included
+        for name, buffer in self.ema_buffers.items():
+            buffer_name = name.split('.')[-1]  # Get the base name
+            # Always include critical buffers regardless of only_save_diff
+            if not self.only_save_diff or buffer_name in self._ALWAYS_INCLUDE_BUFFERS:
+                state_dict[name] = buffer.data
+        
+        # Add internal state
+        state_dict["initted"] = self.initted
+        state_dict["step"] = self.step
+        
         return state_dict
 
     def load_state_dict(self, state_dict):
-        """Load state dict into EMA model."""
-        # Load parameters based on only_save_diff flag
-        for name, param in self.ema_model.named_parameters():
-            if (not self.only_save_diff or param.requires_grad) and name in state_dict:
-                param.data.copy_(state_dict[name].data)
-
-        # Load buffers
-        for name, buffer in self.ema_model.named_buffers():
-            if name in state_dict:
-                buffer.data.copy_(state_dict[name].data)
-
-        # Load internal state
-        if "initted" in state_dict:
-            self.initted.data.copy_(state_dict["initted"].data)
-        if "step" in state_dict:
-            self.step.data.copy_(state_dict["step"].data)
+        """Load state dict with EMA parameters."""
+        for name, param in state_dict.items():
+            if name == "initted":
+                self.initted.data.copy_(param)
+            elif name == "step":
+                self.step.data.copy_(param)
+            elif name in self.ema_params:
+                self.ema_params[name].data.copy_(param)
+            elif name in self.ema_buffers:
+                self.ema_buffers[name].data.copy_(param)