Upload 137 files

hugging/td_fuse/heal.py

@@ -29,6 +29,22 @@ from datasets import load_dataset
 from .config import MergeConfig
 
 
+def _load_model_smart(checkpoint, **kwargs):
+    """Load model — auto-detects Qwen3-VL and uses the correct class."""
+    from transformers import AutoConfig
+    try:
+        config = AutoConfig.from_pretrained(checkpoint, trust_remote_code=True)
+        model_type = getattr(config, 'model_type', '')
+        config_class = type(config).__name__.lower()
+        if 'qwen3_vl' in model_type or 'qwen3vl' in config_class:
+            from transformers import Qwen3VLForConditionalGeneration
+            print(f'[heal] Loading as Qwen3-VL model: {checkpoint}')
+            return Qwen3VLForConditionalGeneration.from_pretrained(checkpoint, **kwargs)
+    except Exception as e:
+        print(f'[heal] Auto-detect failed ({e}), using AutoModelForCausalLM')
+    return AutoModelForCausalLM.from_pretrained(checkpoint, **kwargs)
+
+
 def check_unsloth_available() -> bool:
     """Check if Unsloth is installed and working."""
     try:
@@ -235,7 +251,7 @@ def apply_qlora_standard(
     )
 
     tokenizer = AutoTokenizer.from_pretrained(model_path)
-    model = AutoModelForCausalLM.from_pretrained(
+    model = _load_model_smart(
         model_path,
         quantization_config=bnb_config,
         device_map="auto",