yeq6x commited on
Commit
0d63ccc
·
1 Parent(s): 97617c7

Update app.py to enable GPU support and streamline model download process

Browse files

Modify the run_training function to utilize GPU support with the @spaces.GPU decorator. Additionally, ensure the _startup_download_models function is called during the application startup, enhancing the model loading process. These changes improve performance and user experience by leveraging GPU capabilities and ensuring necessary models are downloaded at launch.

Files changed (1) hide show
  1. app.py +3 -3
app.py CHANGED
@@ -12,7 +12,7 @@ import json
12
 
13
  import gradio as gr
14
  import importlib
15
- # import spaces
16
 
17
  # Local modules
18
  from download_qwen_image_models import download_all_models, DEFAULT_MODELS_DIR
@@ -837,7 +837,7 @@ def _startup_install_musubi_deps() -> None:
837
  print("[QIE] WARN: musubi-tuner installation failed. Continuing.")
838
 
839
 
840
- # @spaces.GPU
841
  def run_training(
842
  output_name: str,
843
  caption: str,
@@ -2008,7 +2008,7 @@ if __name__ == "__main__":
2008
  _startup_install_musubi_deps()
2009
 
2010
  # 2) Download models at startup (blocking by design)
2011
- # _startup_download_models()
2012
 
2013
  # 3) Launch Gradio app
2014
  ui = build_ui()
 
12
 
13
  import gradio as gr
14
  import importlib
15
+ import spaces
16
 
17
  # Local modules
18
  from download_qwen_image_models import download_all_models, DEFAULT_MODELS_DIR
 
837
  print("[QIE] WARN: musubi-tuner installation failed. Continuing.")
838
 
839
 
840
+ @spaces.GPU
841
  def run_training(
842
  output_name: str,
843
  caption: str,
 
2008
  _startup_install_musubi_deps()
2009
 
2010
  # 2) Download models at startup (blocking by design)
2011
+ _startup_download_models()
2012
 
2013
  # 3) Launch Gradio app
2014
  ui = build_ui()