Create convert_weights.py

convert_weights.py

@@ -0,0 +1,70 @@
+from transformers import AutoModel, AutoConfig
+from DaViT.modeling_davit import DaViTModel
+from DaViT.configuration_davit import DaViTConfig
+from unittest.mock import patch
+import os
+import logging
+import requests
+from PIL import Image
+import torch
+from transformers import AutoProcessor, AutoModelForCausalLM
+from unittest.mock import patch
+from transformers.dynamic_module_utils import get_imports
+from typing import Tuple, Dict, Any, Union, List
+
+
+def fixed_get_imports(filename: str | os.PathLike) -> list[str]:
+    """
+    Custom workaround for the import error related to flash_attn.
+    Args:
+        filename (str | os.PathLike): The filename to check for imports.
+    Returns:
+        list[str]: List of required imports.
+    """
+    if not str(filename).endswith("modeling_florence2.py"):
+        return get_imports(filename)
+    imports = get_imports(filename)
+    if "flash_attn" in imports:
+        imports.remove("flash_attn")
+    return imports
+
+
+current_directory = os.getcwd()
+
+# Register the configuration and model
+AutoConfig.register("davit", DaViTConfig)
+AutoModel.register(DaViTConfig, DaViTModel)
+
+
+# Register Huggingface Model
+DaViTConfig.register_for_auto_class()
+DaViTModel.register_for_auto_class("AutoModel")
+
+AutoConfig.register("davit", DaViTConfig)
+AutoModel.register(DaViTConfig, DaViTModel)
+
+# Step 1: Create a configuration object
+config = DaViTConfig()
+with patch("transformers.dynamic_module_utils.get_imports", fixed_get_imports):
+    model = AutoModelForCausalLM.from_pretrained(
+        "microsoft/Florence-2-large-ft",
+        trust_remote_code=True,
+        cache_dir=current_directory,
+        device_map="cpu",
+        torch_dtype=torch.float16,
+    )
+processor = AutoProcessor.from_pretrained(
+    "microsoft/Florence-2-large-ft",
+    trust_remote_code=True,
+    cache_dir=current_directory,
+    device_map="cpu",
+)
+# Step 2: Create a model object
+model2 = AutoModel.from_config(config)
+model2.to(torch.float16)
+
+model2.load_state_dict(model.vision_tower.state_dict())
+
+
+model2.push_to_hub("DaViT-Florence-2-large-ft")
+processor.push_to_hub("DaViT-Florence-2-large-ft")