Revert "Florence2ModelLoader loadmodel"

This reverts commit 8e59d75004da26d5e68c7a9beddfd7f891b26515.

custom_nodes/comfyui-florence2/nodes.py

@@ -198,65 +198,6 @@ 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")
@@ -282,50 +223,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