95 lines
4.0 KiB
Python
95 lines
4.0 KiB
Python
import os
|
|
import sys
|
|
import torch
|
|
import warnings
|
|
import comfy.utils
|
|
|
|
def patch_flex_attention():
|
|
"""
|
|
Patches the torch.nn.attention.flex_attention module to ensure compatibility
|
|
with different PyTorch versions, especially handling the BlockMask renaming
|
|
in PyTorch 2.7.0 nightlies.
|
|
"""
|
|
try:
|
|
import torch.nn.attention.flex_attention as flex_attn
|
|
|
|
# In PyTorch 2.7.0+, BlockMask might be removed or renamed to _BlockMask
|
|
if not hasattr(flex_attn, 'BlockMask'):
|
|
print("[Falcon Perception] Patching flex_attention for PyTorch 2.7.0+ compatibility")
|
|
if hasattr(flex_attn, '_BlockMask'):
|
|
flex_attn.BlockMask = flex_attn._BlockMask
|
|
else:
|
|
class DummyBlockMask:
|
|
pass
|
|
flex_attn.BlockMask = DummyBlockMask
|
|
warnings.warn("BlockMask not found in flex_attention. Creating a dummy class. Model execution might fail.")
|
|
|
|
if not hasattr(flex_attn, 'AuxRequest'):
|
|
if hasattr(flex_attn, '_AuxRequest'):
|
|
flex_attn.AuxRequest = flex_attn._AuxRequest
|
|
else:
|
|
class DummyAuxRequest:
|
|
def __init__(self, **kwargs):
|
|
self.kwargs = kwargs
|
|
flex_attn.AuxRequest = DummyAuxRequest
|
|
|
|
except ImportError:
|
|
warnings.warn("torch.nn.attention.flex_attention is not available in this PyTorch version. Falcon-Perception requires PyTorch >= 2.5.0.")
|
|
|
|
def load_falcon_perception_local(model_path, device_map):
|
|
"""
|
|
Loads the Falcon Perception architecture natively from the bundled node package
|
|
and only loads the .safetensors weights from the user-provided path.
|
|
This completely bypasses all HuggingFace caching and remote code trust issues.
|
|
"""
|
|
patch_flex_attention()
|
|
|
|
# 1. Add the bundled falcon_core directory to sys.path
|
|
core_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "falcon_core")
|
|
if core_dir not in sys.path:
|
|
sys.path.insert(0, core_dir)
|
|
|
|
try:
|
|
# Import the model configuration and class directly from the bundled local files
|
|
from configuration_falcon_perception import FalconPerceptionConfig
|
|
from modeling_falcon_perception import FalconPerceptionForSegmentation
|
|
|
|
# Load configuration natively from the bundled directory
|
|
print(f"[Falcon Perception] Initializing architecture using bundled configs...")
|
|
config = FalconPerceptionConfig.from_pretrained(core_dir)
|
|
|
|
# Initialize empty model structure on CPU (it's small, 0.6B)
|
|
model = FalconPerceptionForSegmentation(config)
|
|
|
|
# Determine weight path
|
|
if os.path.isfile(model_path) and model_path.endswith(".safetensors"):
|
|
weight_path = model_path
|
|
else:
|
|
weight_path = os.path.join(model_path, "model.safetensors")
|
|
if not os.path.exists(weight_path):
|
|
raise FileNotFoundError(f"Could not find model.safetensors inside {model_path} nor as a direct file.")
|
|
|
|
print(f"[Falcon Perception] Loading safetensor weights directly from: {weight_path}")
|
|
sd = comfy.utils.load_torch_file(weight_path, safe_load=True)
|
|
|
|
# Load state dict strictly, but allow some flexibility if keys slightly mismatch
|
|
model.load_state_dict(sd, strict=False)
|
|
print(f"[Falcon Perception] Weights successfully loaded into model structure.")
|
|
|
|
# Move to target device
|
|
target_device = device_map.get("", torch.device("cpu"))
|
|
if target_device != "auto":
|
|
model.to(target_device)
|
|
|
|
model.eval()
|
|
return model
|
|
|
|
except Exception as e:
|
|
import traceback
|
|
traceback.print_exc()
|
|
raise ImportError(f"Failed to initialize self-contained model. Error: {e}")
|
|
finally:
|
|
# Cleanup sys.path
|
|
if sys.path[0] == core_dir:
|
|
sys.path.pop(0)
|