@@ -192,7 +192,7 @@ def put(self, key: str, model: AnyModel) -> None:
192192 self ._cached_models [key ] = cache_record
193193 self ._cache_stack .append (key )
194194 self ._logger .debug (
195- f"Added model { key } (Type: { model .__class__ .__name__ } , Wrap mode: { wrapped_model .__class__ .__name__ } , Model size: { size / MB :.2f} MB)"
195+ f"Added model { key } (Type: { model .__class__ .__name__ } , Wrap mode: { wrapped_model .__class__ .__name__ } , Model size: { size / MB :.2f} MB)"
196196 )
197197
198198 @synchronized
@@ -303,7 +303,7 @@ def _load_locked_model(self, cache_entry: CacheRecord, working_mem_bytes: Option
303303 # 2. If the model can't fit fully into VRAM, then unload all other models and load as much of the model as
304304 # possible.
305305 vram_bytes_freed = self ._offload_unlocked_models (model_vram_needed , working_mem_bytes )
306- self ._logger .debug (f"Unloaded models (if necessary): vram_bytes_freed={ (vram_bytes_freed / MB ):.2f} MB" )
306+ self ._logger .debug (f"Unloaded models (if necessary): vram_bytes_freed={ (vram_bytes_freed / MB ):.2f} MB" )
307307
308308 # Check the updated vram_available after offloading.
309309 vram_available = self ._get_vram_available (working_mem_bytes )
@@ -317,7 +317,7 @@ def _load_locked_model(self, cache_entry: CacheRecord, working_mem_bytes: Option
317317 vram_bytes_freed_from_own_model = self ._move_model_to_ram (cache_entry , - vram_available )
318318 vram_available = self ._get_vram_available (working_mem_bytes )
319319 self ._logger .debug (
320- f"Unloaded { vram_bytes_freed_from_own_model / MB :.2f} MB from the model being locked ({ cache_entry .key } )."
320+ f"Unloaded { vram_bytes_freed_from_own_model / MB :.2f} MB from the model being locked ({ cache_entry .key } )."
321321 )
322322
323323 # Move as much of the model as possible into VRAM.
@@ -333,10 +333,12 @@ def _load_locked_model(self, cache_entry: CacheRecord, working_mem_bytes: Option
333333 self ._logger .info (
334334 f"Loaded model '{ cache_entry .key } ' ({ cache_entry .cached_model .model .__class__ .__name__ } ) onto "
335335 f"{ self ._execution_device .type } device in { (time .time () - start_time ):.2f} s. "
336- f"Total model size: { model_total_bytes / MB :.2f} MB, "
337- f"VRAM: { model_cur_vram_bytes / MB :.2f} MB ({ loaded_percent :.1%} )"
336+ f"Total model size: { model_total_bytes / MB :.2f} MB, "
337+ f"VRAM: { model_cur_vram_bytes / MB :.2f} MB ({ loaded_percent :.1%} )"
338+ )
339+ self ._logger .debug (
340+ f"Loaded model onto execution device: model_bytes_loaded={ (model_bytes_loaded / MB ):.2f} MB, "
338341 )
339- self ._logger .debug (f"Loaded model onto execution device: model_bytes_loaded={ (model_bytes_loaded / MB ):.2f} MB, " )
340342 self ._logger .debug (
341343 f"After loading: { self ._get_vram_state_str (model_cur_vram_bytes , model_total_bytes , vram_available )} "
342344 )
@@ -495,10 +497,10 @@ def _get_vram_state_str(self, model_cur_vram_bytes: int, model_total_bytes: int,
495497 """Helper function for preparing a VRAM state log string."""
496498 model_cur_vram_bytes_percent = model_cur_vram_bytes / model_total_bytes if model_total_bytes > 0 else 0
497499 return (
498- f"model_total={ model_total_bytes / MB :.0f} MB, "
499- + f"model_vram={ model_cur_vram_bytes / MB :.0f} MB ({ model_cur_vram_bytes_percent :.1%} %), "
500+ f"model_total={ model_total_bytes / MB :.0f} MB, "
501+ + f"model_vram={ model_cur_vram_bytes / MB :.0f} MB ({ model_cur_vram_bytes_percent :.1%} %), "
500502 # + f"vram_total={int(self._max_vram_cache_size * GB)/MB:.0f} MB, "
501- + f"vram_available={ (vram_available / MB ):.0f} MB, "
503+ + f"vram_available={ (vram_available / MB ):.0f} MB, "
502504 )
503505
504506 def _offload_unlocked_models (self , vram_bytes_required : int , working_mem_bytes : Optional [int ] = None ) -> int :
@@ -509,7 +511,7 @@ def _offload_unlocked_models(self, vram_bytes_required: int, working_mem_bytes:
509511 int: The number of bytes freed based on believed model sizes. The actual change in VRAM may be different.
510512 """
511513 self ._logger .debug (
512- f"Offloading unlocked models with goal of making room for { vram_bytes_required / MB :.2f} MB of VRAM."
514+ f"Offloading unlocked models with goal of making room for { vram_bytes_required / MB :.2f} MB of VRAM."
513515 )
514516 vram_bytes_freed = 0
515517 # TODO(ryand): Give more thought to the offloading policy used here.
@@ -527,7 +529,7 @@ def _offload_unlocked_models(self, vram_bytes_required: int, working_mem_bytes:
527529 cache_entry_bytes_freed = self ._move_model_to_ram (cache_entry , vram_bytes_to_free )
528530 if cache_entry_bytes_freed > 0 :
529531 self ._logger .debug (
530- f"Unloaded { cache_entry .key } from VRAM to free { (cache_entry_bytes_freed / MB ):.0f} MB."
532+ f"Unloaded { cache_entry .key } from VRAM to free { (cache_entry_bytes_freed / MB ):.0f} MB."
531533 )
532534 vram_bytes_freed += cache_entry_bytes_freed
533535
@@ -609,7 +611,7 @@ def make_room(self, bytes_needed: int) -> None:
609611 external references to the model, there's nothing that the cache can do about it, and those models will not be
610612 garbage-collected.
611613 """
612- self ._logger .debug (f"Making room for { bytes_needed / MB :.2f} MB of RAM." )
614+ self ._logger .debug (f"Making room for { bytes_needed / MB :.2f} MB of RAM." )
613615 self ._log_cache_state (title = "Before dropping models:" )
614616
615617 ram_bytes_available = self ._get_ram_available ()
@@ -625,7 +627,7 @@ def make_room(self, bytes_needed: int) -> None:
625627 if not cache_entry .is_locked :
626628 ram_bytes_freed += cache_entry .cached_model .total_bytes ()
627629 self ._logger .debug (
628- f"Dropping { model_key } from RAM cache to free { (cache_entry .cached_model .total_bytes ()/ MB ):.2f} MB."
630+ f"Dropping { model_key } from RAM cache to free { (cache_entry .cached_model .total_bytes () / MB ):.2f} MB."
629631 )
630632 self ._delete_cache_entry (cache_entry )
631633 del cache_entry
@@ -650,7 +652,7 @@ def make_room(self, bytes_needed: int) -> None:
650652 gc .collect ()
651653
652654 TorchDevice .empty_cache ()
653- self ._logger .debug (f"Dropped { models_cleared } models to free { ram_bytes_freed / MB :.2f} MB of RAM." )
655+ self ._logger .debug (f"Dropped { models_cleared } models to free { ram_bytes_freed / MB :.2f} MB of RAM." )
654656 self ._log_cache_state (title = "After dropping models:" )
655657
656658 def _delete_cache_entry (self , cache_entry : CacheRecord ) -> None :
0 commit comments