Update app.py

app.py

@@ -7,7 +7,7 @@ from PIL import Image
 import gradio as gr
 
 # ===============================
-# Make dummy packages for flash_attn and ps3 (CPU-friendly import stubs)
+# Helper to create package-like dummy modules
 # ===============================
 def _mk_pkg(name: str):
     m = types.ModuleType(name)
@@ -17,7 +17,18 @@ def _mk_pkg(name: str):
     m.__path__ = []
     return m
 
-# --- flash_attn package + submodules ---
+# ===============================
+# Disable GPU-only/optional paths
+# ===============================
+os.environ.setdefault("FLASH_ATTENTION", "0")
+os.environ.setdefault("XFORMERS_DISABLED", "1")
+os.environ.setdefault("ACCELERATE_USE_DEVICE_MAP", "0")
+os.environ.setdefault("DISABLE_TRITON", "1")   # avoid triton kernels
+# os.environ.setdefault("CUDA_VISIBLE_DEVICES", "")  # uncomment to force CPU
+
+# ===============================
+# flash_attn stubs (package + submodules)
+# ===============================
 flash_attn_pkg = _mk_pkg("flash_attn")
 
 flash_attn_interface = types.ModuleType("flash_attn.flash_attn_interface")
@@ -31,7 +42,6 @@ flash_attn_bert_padding.__spec__ = importlib.machinery.ModuleSpec(
 )
 
 def _dummy_func(*args, **kwargs):
-    # Should never be called on CPU; if it is, fail loudly so we notice.
     raise RuntimeError("flash_attn is not available in this environment.")
 
 flash_attn_interface.flash_attn_unpadded_qkvpacked_func = _dummy_func
@@ -43,7 +53,9 @@ 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
 
-# --- ps3 package stub ---
+# ===============================
+# ps3 stub (optional vision tower)
+# ===============================
 ps3_pkg = _mk_pkg("ps3")
 class _PS3Config: pass
 class _PS3VisionConfig: pass
@@ -56,16 +68,20 @@ ps3_pkg.PS3VisionModel = _PS3VisionModel
 sys.modules["ps3"] = ps3_pkg
 
 # ===============================
-# Runtime env (CPU-safe defaults)
+# Quantization stub to avoid Triton path
+# VILA falls back to "from FloatPointQuantizeTorch import *" if Triton import fails.
+# Provide a tiny no-op module so imports succeed.
 # ===============================
-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 on the Space
-# os.environ.setdefault("CUDA_VISIBLE_DEVICES", "")
+fpqt = types.ModuleType("FloatPointQuantizeTorch")
+def _id(x, *a, **k): return x  # identity
+# names used by llava.model.qfunction
+fpqt.block_cut = _id
+fpqt.block_quant = _id
+fpqt.block_reshape = _id
+sys.modules["FloatPointQuantizeTorch"] = fpqt
 
 # ===============================
-# VILA imports & model load
+# Load VILA
 # ===============================
 from llava.model.builder import load_pretrained_model
 from llava.constants import DEFAULT_IMAGE_TOKEN
@@ -77,7 +93,6 @@ try:
         MODEL_PATH, model_name="", model_base=None
     )
 except Exception as e:
-    # Surface a friendly error on the UI instead of crashing the Space
     ERR = f"Failed to load model '{MODEL_PATH}': {e}"
     def _boot_error_ui():
         with gr.Blocks(title="VILA 1.5 3B – Error") as demo:
@@ -105,7 +120,6 @@ def vila_infer(image, prompt):
 
     pil = Image.fromarray(image).convert("RGB")
 
-    # Minimal multimodal conversation: image + text
     try:
         out = model.generate_content(
             prompt=[{
@@ -115,18 +129,18 @@ def vila_infer(image, prompt):
                     {"type": "text", "value": prompt}
                 ]
             }],
-            generation_config=None  # use model defaults
+            generation_config=None
         )
         return str(out).strip()
     except Exception as e:
         return f"❌ Inference error: {e}"
 
 # ===============================
-# Gradio UI
+# UI
 # ===============================
 with gr.Blocks(title="VILA 1.5 3B (HF Space)") as demo:
     gr.Markdown("## 🖼️ VILA-1.5-3B — Image Description Demo")
-    gr.Markdown("Upload an image and press **Run**. Leave the prompt as default for simple captioning.")
+    gr.Markdown("Upload an image and press **Run**.")
 
     with gr.Row():
         img = gr.Image(type="numpy", label="Image", height=320)
@@ -137,4 +151,4 @@ with gr.Blocks(title="VILA 1.5 3B (HF Space)") as demo:
 
     run_btn.click(vila_infer, [img, prompt], out)
 
-demo.launch()
+demo.launch()