nodes/kijai-ComfyUI-Florence2
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Initial load to offload device instead of device_map

Browse Source
This commit is contained in:
Silver
2025-04-14 12:03:39 +02:00
committed by GitHub
parent 8d71159445
commit 820689f6dd
1 changed files with 3 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nodes.py
+3 -2
View File
@@ -123,7 +123,7 @@ class DownloadAndLoadFlorence2Model:
print(f"Florence2 using {attention} for attention")
with patch("transformers.dynamic_module_utils.get_imports", fixed_get_imports): #workaround for unnecessary flash_attn requirement
model = AutoModelForCausalLM.from_pretrained(model_path, attn_implementation=attention, device_map=device, torch_dtype=dtype,trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(model_path, attn_implementation=attention, torch_dtype=dtype,trust_remote_code=True).to(offload_device)
processor = AutoProcessor.from_pretrained(model_path, trust_remote_code=True)
if lora is not None:
@@ -197,12 +197,13 @@ class Florence2ModelLoader:
def loadmodel(self, model, precision, attention, lora=None):
device = mm.get_torch_device()
offload_device = mm.unet_offload_device()
dtype = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}[precision]
model_path = Florence2ModelLoader.model_paths.get(model)
print(f"Loading model from {model_path}")
print(f"Florence2 using {attention} for attention")
with patch("transformers.dynamic_module_utils.get_imports", fixed_get_imports): #workaround for unnecessary flash_attn requirement
model = AutoModelForCausalLM.from_pretrained(model_path, attn_implementation=attention, device_map=device, torch_dtype=dtype,trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(model_path, attn_implementation=attention, torch_dtype=dtype,trust_remote_code=True).to(offload_device)
processor = AutoProcessor.from_pretrained(model_path, trust_remote_code=True)
if lora is not None: