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
@@ -18,9 +25,22 @@ from modeling.bagel import (
 )
 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" #Download from https://huggingface.co/ByteDance-Seed/BAGEL-7B-MoT
 
 llm_config = Qwen2Config.from_json_file(os.path.join(model_path, "llm_config.json"))
 llm_config.qk_norm = True
@@ -300,6 +320,7 @@ with gr.Blocks() as demo:
         )
         
         # Process function based on thinking option and hyperparameters
+        @spaces.GPU(duration=90)
         def process_text_to_image(prompt, show_thinking, cfg_text_scale, 
                                  cfg_interval, timestep_shift, 
                                  num_timesteps, cfg_renorm_min, cfg_renorm_type, 
@@ -391,6 +412,7 @@ with gr.Blocks() as demo:
         )
         
         # Process editing with thinking option and hyperparameters
+        @spaces.GPU(duration=90)
         def process_edit_image(image, prompt, show_thinking, cfg_text_scale, 
                               cfg_img_scale, cfg_interval, 
                               timestep_shift, num_timesteps, cfg_renorm_min, 
@@ -444,6 +466,7 @@ with gr.Blocks() as demo:
         img_understand_btn = gr.Button("Submit")
         
         # Process understanding with thinking option and hyperparameters
+        @spaces.GPU(duration=90)
         def process_understanding(image, prompt, show_thinking, do_sample, 
                                  text_temperature, max_new_tokens):
             result = image_understanding(
@@ -502,4 +525,4 @@ with gr.Blocks() as demo:
 </div>
 """)
 
-demo.launch(share=True)
+demo.launch()