Update app.py

app.py

@@ -1,8 +1,15 @@
+import spaces
 import gradio as gr
 import numpy as np
 import os
 import torch
 import random
+import subprocess
+subprocess.run(
+    "pip install flash-attn --no-build-isolation",
+    env={"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"},
+    shell=True,
+)
 
 from accelerate import infer_auto_device_map, load_checkpoint_and_dispatch, init_empty_weights
 from PIL import Image
@@ -17,10 +24,23 @@ from modeling.bagel import (
     SiglipVisionConfig, SiglipVisionModel
 )
 from modeling.qwen2 import Qwen2Tokenizer
+from huggingface_hub import snapshot_download
+
+save_dir = "./model"
+repo_id = "ByteDance-Seed/BAGEL-7B-MoT"
+cache_dir = save_dir + "/cache"
+
+snapshot_download(cache_dir=cache_dir,
+  local_dir=save_dir,
+  repo_id=repo_id,
+  local_dir_use_symlinks=False,
+  resume_download=True,
+  allow_patterns=["*.json", "*.safetensors", "*.bin", "*.py", "*.md", "*.txt"],
+)
 
 
 # Model Initialization
-model_path = "/path/to/BAGEL-7B-MoT/weights" #Download from https://huggingface.co/ByteDance-Seed/BAGEL-7B-MoT
+model_path = "./model"
 
 llm_config = Qwen2Config.from_json_file(os.path.join(model_path, "llm_config.json"))
 llm_config.qk_norm = True
@@ -502,4 +522,4 @@ with gr.Blocks() as demo:
 </div>
 """)
 
-demo.launch(share=True)
+demo.launch()