Update app.py

app.py

@@ -6,43 +6,69 @@ from PIL import Image
 import gradio as gr
 
 # ===============================
-# Patch flash_attn for CPU runtime
+# Make a PACKAGE-like dummy flash_attn
 # ===============================
-dummy_flash_attn = types.ModuleType("flash_attn")
-dummy_flash_attn.__spec__ = importlib.machinery.ModuleSpec("flash_attn", loader=None)
+def _mk_pkg(name: str):
+    m = types.ModuleType(name)
+    # Mark as a package: give it a spec with submodule locations and a __path__
+    spec = importlib.machinery.ModuleSpec(name, loader=None, is_package=True)
+    spec.submodule_search_locations = []  # important: tells importlib it's a package
+    m.__spec__ = spec
+    m.__path__ = []  # also marks as package
+    return m
 
-dummy_interface = types.ModuleType("flash_attn.flash_attn_interface")
-dummy_interface.__spec__ = importlib.machinery.ModuleSpec(
+# Root package
+flash_attn_pkg = _mk_pkg("flash_attn")
+
+# Submodule: flash_attn.flash_attn_interface
+flash_attn_interface = types.ModuleType("flash_attn.flash_attn_interface")
+flash_attn_interface.__spec__ = importlib.machinery.ModuleSpec(
     "flash_attn.flash_attn_interface", loader=None
 )
 
+# Submodule: flash_attn.bert_padding
+flash_attn_bert_padding = types.ModuleType("flash_attn.bert_padding")
+flash_attn_bert_padding.__spec__ = importlib.machinery.ModuleSpec(
+    "flash_attn.bert_padding", loader=None
+)
+
 def _dummy_func(*args, **kwargs):
+    # Should never be called on CPU; if it is, let’s fail loudly
     raise RuntimeError("flash_attn is not available in this environment.")
 
-dummy_interface.flash_attn_unpadded_qkvpacked_func = _dummy_func
-dummy_interface.flash_attn_varlen_qkvpacked_func = _dummy_func
+# Functions some imports expect to exist:
+flash_attn_interface.flash_attn_unpadded_qkvpacked_func = _dummy_func
+flash_attn_interface.flash_attn_varlen_qkvpacked_func = _dummy_func
+flash_attn_bert_padding.pad_input = _dummy_func
+flash_attn_bert_padding.unpad_input = _dummy_func
 
-sys.modules["flash_attn"] = dummy_flash_attn
-sys.modules["flash_attn.flash_attn_interface"] = dummy_interface
+# Register modules
+sys.modules["flash_attn"] = flash_attn_pkg
+sys.modules["flash_attn.flash_attn_interface"] = flash_attn_interface
+sys.modules["flash_attn.bert_padding"] = flash_attn_bert_padding
 
 # ===============================
-# Hugging Face model setup
+# Runtime env (CPU-friendly)
 # ===============================
 os.environ.setdefault("FLASH_ATTENTION", "0")
 os.environ.setdefault("XFORMERS_DISABLED", "1")
 os.environ.setdefault("ACCELERATE_USE_DEVICE_MAP", "0")
+# Uncomment to force CPU even if a GPU is present:
+# os.environ.setdefault("CUDA_VISIBLE_DEVICES", "")
 
+# ===============================
+# VILA imports & load
+# ===============================
 from llava.model.builder import load_pretrained_model
 from llava.constants import DEFAULT_IMAGE_TOKEN
 
 MODEL_PATH = "Efficient-Large-Model/VILA1.5-3b"
 
-# Load model + tokenizer + image processor
 tokenizer, model, image_processor, context_len = load_pretrained_model(
     MODEL_PATH, model_name="", model_base=None
 )
 
-# Add a fallback chat template
+# Fallback chat template if missing
 if getattr(tokenizer, "chat_template", None) is None:
     tokenizer.chat_template = (
         "{% for message in messages %}{{ message['role'] | upper }}: "