Florence2ModelLoader loadmodel

custom_nodes/comfyui-florence2/nodes.py

@@ -198,6 +198,65 @@ class DownloadAndLoadFlorence2Lora:
     
 class Florence2ModelLoader:
 
+    # 用下面的函数完整替换掉旧的 loadmodel 函数
+    def loadmodel(self, model, precision, attention, lora=None, convert_to_safetensors=False):
+        """
+        一个为 Hugging Face Spaces ZeroGPU 环境重写的、稳健的 loadmodel 函数。
+        它完全移除了手动的设备管理，并使用 accelerate 库进行智能调度。
+        """
+        # 1. 彻底删除所有手动的设备管理
+        # device = mm.get_torch_device()  <-- 已删除
+        # offload_device = mm.unet_offload_device()  <-- 已删除
+        
+        dtype = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}[precision]
+
+        model_path = self.model_paths.get(model)
+        print(f"Loading model from {model_path} using the correct Spaces method (device_map='auto').")
+
+        # 2. 保留 safetensors 转换逻辑，但修复 map_location
+        if convert_to_safetensors:
+            model_weight_path = os.path.join(model_path, 'pytorch_model.bin')
+            safetensors_weight_path = os.path.join(model_path, 'model.safetensors')
+            if os.path.exists(model_weight_path) and not os.path.exists(safetensors_weight_path):
+                print(f"Converting {model_weight_path} to {safetensors_weight_path}")
+                # 使用 "cpu" 作为 map_location 确保在任何环境下都安全
+                sd = torch.load(model_weight_path, map_location="cpu")
+                save_file(sd, safetensors_weight_path)
+                if os.path.exists(safetensors_weight_path):
+                    os.remove(model_weight_path)
+                    print(f"Conversion successful. Original file deleted.")
+
+        # 3. 统一使用 from_pretrained 和 device_map="auto" 加载模型
+        #    删除所有 .to(device) 调用
+        
+        # 假设 Florence2ForConditionalGeneration 是你的主要模型类
+        from .modeling_florence2 import Florence2ForConditionalGeneration
+        
+        print("Loading model with device_map='auto'...")
+        model_instance = Florence2ForConditionalGeneration.from_pretrained(
+            model_path, 
+            attn_implementation=attention, 
+            torch_dtype=dtype,
+            device_map="auto",
+            low_cpu_mem_usage=True # 强烈推荐，防止CPU内存溢出
+        )
+        print("Model loaded successfully onto meta device / CPU.")
+            
+        processor = AutoProcessor.from_pretrained(model_path, trust_remote_code=True)
+
+        if lora is not None:
+            from peft import PeftModel
+            # PEFT 会自动处理设备，无需改动
+            model_instance = PeftModel.from_pretrained(model_instance, lora, trust_remote_code=True)
+        
+        florence2_model = {
+            'model': model_instance, 
+            'processor': processor,
+            'dtype': dtype
+        }
+   
+        return (florence2_model,) # 保持返回元组的格式
+        
     @classmethod
     def INPUT_TYPES(s):
         all_llm_paths = folder_paths.get_folder_paths("LLM")
@@ -223,50 +282,50 @@ class Florence2ModelLoader:
     FUNCTION = "loadmodel"
     CATEGORY = "Florence2"
 
-    def loadmodel(self, model, precision, attention, lora=None, convert_to_safetensors=False):
-        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")
-        if convert_to_safetensors:
-            model_weight_path = os.path.join(model_path, 'pytorch_model.bin')
-            if os.path.exists(model_weight_path):
-                safetensors_weight_path = os.path.join(model_path, 'model.safetensors')
-                print(f"Converting {model_weight_path} to {safetensors_weight_path}")
-                if not os.path.exists(safetensors_weight_path):
-                    sd = torch.load(model_weight_path, map_location=offload_device)
-                    sd_new = {}
-                    for k, v in sd.items():
-                        sd_new[k] = v.clone()
-                    save_file(sd_new, safetensors_weight_path)
-                    if os.path.exists(safetensors_weight_path):
-                        print(f"Conversion successful. Deleting original file: {model_weight_path}")
-                        os.remove(model_weight_path)
-                        print(f"Original {model_weight_path} file deleted.")
-
-        if transformers.__version__ < '4.51.0':
-            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, torch_dtype=dtype,trust_remote_code=True).to(offload_device)
-        else:
-            from .modeling_florence2 import Florence2ForConditionalGeneration
-            model = Florence2ForConditionalGeneration.from_pretrained(model_path, attn_implementation=attention, torch_dtype=dtype).to(offload_device)
-        processor = AutoProcessor.from_pretrained(model_path, trust_remote_code=True)
-
-        if lora is not None:
-            from peft import PeftModel
-            adapter_name = lora
-            model = PeftModel.from_pretrained(model, adapter_name, trust_remote_code=True)
+    # def loadmodel(self, model, precision, attention, lora=None, convert_to_safetensors=False):
+    #     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")
+    #     if convert_to_safetensors:
+    #         model_weight_path = os.path.join(model_path, 'pytorch_model.bin')
+    #         if os.path.exists(model_weight_path):
+    #             safetensors_weight_path = os.path.join(model_path, 'model.safetensors')
+    #             print(f"Converting {model_weight_path} to {safetensors_weight_path}")
+    #             if not os.path.exists(safetensors_weight_path):
+    #                 sd = torch.load(model_weight_path, map_location=offload_device)
+    #                 sd_new = {}
+    #                 for k, v in sd.items():
+    #                     sd_new[k] = v.clone()
+    #                 save_file(sd_new, safetensors_weight_path)
+    #                 if os.path.exists(safetensors_weight_path):
+    #                     print(f"Conversion successful. Deleting original file: {model_weight_path}")
+    #                     os.remove(model_weight_path)
+    #                     print(f"Original {model_weight_path} file deleted.")
+
+    #     if transformers.__version__ < '4.51.0':
+    #         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, torch_dtype=dtype,trust_remote_code=True).to(offload_device)
+    #     else:
+    #         from .modeling_florence2 import Florence2ForConditionalGeneration
+    #         model = Florence2ForConditionalGeneration.from_pretrained(model_path, attn_implementation=attention, torch_dtype=dtype).to(offload_device)
+    #     processor = AutoProcessor.from_pretrained(model_path, trust_remote_code=True)
+
+    #     if lora is not None:
+    #         from peft import PeftModel
+    #         adapter_name = lora
+    #         model = PeftModel.from_pretrained(model, adapter_name, trust_remote_code=True)
         
-        florence2_model = {
-            'model': model, 
-            'processor': processor,
-            'dtype': dtype
-            }
+    #     florence2_model = {
+    #         'model': model, 
+    #         'processor': processor,
+    #         'dtype': dtype
+    #         }
    
-        return (florence2_model,)
+    #     return (florence2_model,)
     
 class Florence2Run:
     @classmethod