1580158
/
automatic
mirror of https://github.com/vladmandic/automatic
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

refactor(caption): remove unused _device tracking property

pull/4569/head
CalamitousFelicitousness 2026-01-21 10:51:40 +00:00
parent 6b10f0df4f
commit 26c679f9e7
2 changed files with 6 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
cli/test-tagger.py
View File

@ -271,19 +271,19 @@ class TaggerTest:
print(f" Baseline: GPU={baseline['gpu_allocated']:.1f}MB, RAM={baseline['ram_used']:.1f}MB")
# Test 1: Check initial state (should be on CPU after load)
initial_device = deepbooru.model._device
initial_device = next(deepbooru.model.model.parameters()).device
print(f" Initial device: {initial_device}")
if initial_device == devices.cpu:
if initial_device.type == 'cpu':
self.log_pass("DeepBooru: initial state on CPU")
else:
self.log_pass(f"DeepBooru: initial state on {initial_device}")
# Test 2: Move to GPU (start)
deepbooru.model.start()
gpu_device = deepbooru.model._device
gpu_device = next(deepbooru.model.model.parameters()).device
after_gpu = self.get_memory_stats()
print(f" After start(): {gpu_device} | GPU={after_gpu['gpu_allocated']:.1f}MB (+{after_gpu['gpu_allocated']-baseline['gpu_allocated']:.1f}MB)")
if gpu_device == devices.device:
if gpu_device.type == devices.device.type:
self.log_pass(f"DeepBooru: moved to GPU ({gpu_device})")
else:
self.log_fail(f"DeepBooru: failed to move to GPU, got {gpu_device}")
@ -306,9 +306,9 @@ class TaggerTest:
if torch.cuda.is_available():
torch.cuda.empty_cache()
after_offload = self.get_memory_stats()
cpu_device = deepbooru.model._device
cpu_device = next(deepbooru.model.model.parameters()).device
print(f" After stop(): {cpu_device} | GPU={after_offload['gpu_allocated']:.1f}MB, RAM={after_offload['ram_used']:.1f}MB")
if cpu_device == devices.cpu:
if cpu_device.type == 'cpu':
self.log_pass("DeepBooru: offloaded to CPU")
else:
self.log_fail(f"DeepBooru: failed to offload, still on {cpu_device}")

5
modules/interrogate/deepbooru.py
View File

@ -13,7 +13,6 @@ load_lock = threading.Lock()
class DeepDanbooru:
def __init__(self):
self.model = None
self._device = devices.cpu
def load(self):
with load_lock:
@ -33,17 +32,14 @@ class DeepDanbooru:
self.model.load_state_dict(torch.load(files[0], map_location="cpu"))
self.model.eval()
self.model.to(devices.cpu, devices.dtype)
self._device = devices.cpu
def start(self):
self.load()
self.model.to(devices.device)
self._device = devices.device
def stop(self):
if shared.opts.interrogate_offload:
self.model.to(devices.cpu)
self._device = devices.cpu
devices.torch_gc()
def tag(self, pil_image, **kwargs):
@ -175,7 +171,6 @@ def unload_model():
if model.model is not None:
shared.log.debug('DeepBooru unload')
model.model = None
model._device = devices.cpu
devices.torch_gc(force=True)