| 1132 | return output.movedim(1, -1) |
| 1133 | |
| 1134 | def encode(self, pixel_samples): |
| 1135 | self.throw_exception_if_invalid() |
| 1136 | pixel_samples = self.vae_encode_crop_pixels(pixel_samples) |
| 1137 | pixel_samples = pixel_samples.movedim(-1, 1) |
| 1138 | do_tile = False |
| 1139 | if self.latent_dim == 3 and pixel_samples.ndim < 5: |
| 1140 | if not self.not_video: |
| 1141 | pixel_samples = pixel_samples.movedim(1, 0).unsqueeze(0) |
| 1142 | else: |
| 1143 | pixel_samples = pixel_samples.unsqueeze(2) |
| 1144 | |
| 1145 | with model_management.cuda_device_context(self.device): |
| 1146 | try: |
| 1147 | memory_used = self.memory_used_encode(pixel_samples.shape, self.vae_dtype) |
| 1148 | model_management.load_models_gpu([self.patcher], memory_required=memory_used, force_full_load=self.disable_offload) |
| 1149 | free_memory = self.patcher.get_free_memory(self.device) |
| 1150 | batch_number = int(free_memory / max(1, memory_used)) |
| 1151 | batch_number = max(1, batch_number) |
| 1152 | samples = None |
| 1153 | for x in range(0, pixel_samples.shape[0], batch_number): |
| 1154 | pixels_in = self.process_input(pixel_samples[x:x + batch_number]).to(self.vae_dtype) |
| 1155 | if getattr(self.first_stage_model, 'comfy_has_chunked_io', False): |
| 1156 | out = self.first_stage_model.encode(pixels_in, device=self.device) |
| 1157 | else: |
| 1158 | pixels_in = pixels_in.to(self.device) |
| 1159 | out = self.first_stage_model.encode(pixels_in) |
| 1160 | out = out.to(self.output_device).to(dtype=self.vae_output_dtype()) |
| 1161 | if samples is None: |
| 1162 | samples = torch.empty((pixel_samples.shape[0],) + tuple(out.shape[1:]), device=self.output_device, dtype=self.vae_output_dtype()) |
| 1163 | samples[x:x + batch_number] = out |
| 1164 | |
| 1165 | except Exception as e: |
| 1166 | model_management.raise_non_oom(e) |
| 1167 | logging.warning("Warning: Ran out of memory when regular VAE encoding, retrying with tiled VAE encoding.") |
| 1168 | #NOTE: We don't know what tensors were allocated to stack variables at the time of the |
| 1169 | #exception and the exception itself refs them all until we get out of this except block. |
| 1170 | #So we just set a flag for tiler fallback so that tensor gc can happen once the |
| 1171 | #exception is fully off the books. |
| 1172 | do_tile = True |
| 1173 | |
| 1174 | if do_tile: |
| 1175 | comfy.model_management.soft_empty_cache() |
| 1176 | if self.latent_dim == 3: |
| 1177 | tile = 256 |
| 1178 | overlap = tile // 4 |
| 1179 | samples = self.encode_tiled_3d(pixel_samples, tile_x=tile, tile_y=tile, overlap=(1, overlap, overlap)) |
| 1180 | elif self.latent_dim == 1 or self.extra_1d_channel is not None: |
| 1181 | samples = self.encode_tiled_1d(pixel_samples) |
| 1182 | else: |
| 1183 | samples = self.encode_tiled_(pixel_samples) |
| 1184 | |
| 1185 | return samples |
| 1186 | |
| 1187 | def encode_tiled(self, pixel_samples, tile_x=None, tile_y=None, overlap=None, tile_t=None, overlap_t=None): |
| 1188 | self.throw_exception_if_invalid() |