Enhance training functionality with ZeroGPU support and UI adjustments. Added options to override max epochs and save frequency, and implemented GPU request handling for Spaces compatibility.

app.py

@@ -9,9 +9,10 @@ from pathlib import Path
 from typing import Dict, Iterable, List, Optional
 
 import gradio as gr
-
 import spaces
 
+# No Spaces GPU reservation to allow zero-GPU (CPU-only) usage
+
 # Local modules
 from download_qwen_image_models import download_all_models, DEFAULT_MODELS_DIR
 
@@ -73,6 +74,8 @@ def _prepare_script(
     models_root: str,
     output_dir_base: Optional[str] = None,
     dataset_config: Optional[str] = None,
+    override_max_epochs: Optional[int] = None,
+    override_save_every: Optional[int] = None,
 ) -> Path:
     """Create a temporary copy of train_QIE.sh with injected variables.
 
@@ -137,6 +140,24 @@ def _prepare_script(
     txt = _replace_model_path(txt, "text_encoder", "text_encoder/qwen_2.5_vl_7b.safetensors")
     txt = _replace_model_path(txt, "dit", "dit/qwen_image_edit_2509_bf16.safetensors")
 
+    # ZeroGPU compatibility: avoid spawning via 'accelerate launch'.
+    # Run the training module directly in-process so GPU stays attached
+    # to the same Python request context.
+    txt = re.sub(
+        r"\baccelerate\s+launch\s+src/musubi_tuner/qwen_image_train_network.py",
+        r"python src/musubi_tuner/qwen_image_train_network.py",
+        txt,
+        flags=re.MULTILINE,
+    )
+
+    # Optionally override epochs and save frequency for ZeroGPU time slicing
+    if override_max_epochs is not None and override_max_epochs > 0:
+        txt = re.sub(r"--max_train_epochs\s+\d+",
+                     f"--max_train_epochs {override_max_epochs}", txt)
+    if override_save_every is not None and override_save_every > 0:
+        txt = re.sub(r"--save_every_n_epochs\s+\d+",
+                     f"--save_every_n_epochs {override_save_every}", txt)
+
     # Write to a temp file alongside this repo for easier inspection
     run_dir = TRAINING_DIR / ".gradio_runs"
     run_dir.mkdir(parents=True, exist_ok=True)
@@ -199,6 +220,7 @@ def _startup_clone_musubi_tuner() -> None:
         print(f"[QIE] Clone failed: {e}")
 
 
+@spaces.GPU(duration=7200)
 def run_training(
     dataset_name: str,
     caption: str,
@@ -215,6 +237,8 @@ def run_training(
     models_root: str,
     output_dir_base: str,
     dataset_config: str,
+    max_epochs: int,
+    save_every: int,
 ) -> Iterable[str]:
     # Basic validation
     if not dataset_name.strip():
@@ -241,8 +265,11 @@ def run_training(
         models_root=models_root.strip() or DEFAULT_MODELS_ROOT,
         output_dir_base=(output_dir_base.strip() or None),
         dataset_config=(dataset_config.strip() or None),
+        override_max_epochs=max_epochs if max_epochs and max_epochs > 0 else None,
+        override_save_every=save_every if save_every and save_every > 0 else None,
     )
 
+
     shell = _pick_shell()
     yield f"[QIE] Using shell: {shell}"
     yield f"[QIE] Running script: {tmp_script}"
@@ -300,12 +327,17 @@ def build_ui() -> gr.Blocks:
         run_btn = gr.Button("Start Training", variant="primary")
         logs = gr.Textbox(label="Logs", lines=20)
 
+        with gr.Row():
+            max_epochs = gr.Number(label="Max epochs (this run)", value=10, precision=0)
+            save_every = gr.Number(label="Save every N epochs", value=5, precision=0)
+
         run_btn.click(
             fn=run_training,
             inputs=[
                 dataset_name, caption, data_root, image_folder,
                 c0, c1, c2, c3, c4, c5, c6, c7,
                 models_root, output_dir_base, dataset_config,
+                max_epochs, save_every,
             ],
             outputs=logs,
         )
@@ -313,6 +345,11 @@ def build_ui() -> gr.Blocks:
     return demo
 
 
+@spaces.GPU(duration=600)
+def _request_gpu_on_startup() -> str:
+    return "gpu-requested"
+
+
 def _startup_download_models() -> None:
     models_dir = DEFAULT_MODELS_ROOT
     print(f"[QIE] Ensuring models in: {models_dir}")
@@ -323,6 +360,13 @@ def _startup_download_models() -> None:
 
 
 if __name__ == "__main__":
+    # 0) Request GPU immediately for Spaces dynamic hardware
+    try:
+        tag = _request_gpu_on_startup()
+        print(f"[QIE] Spaces GPU tag: {tag}")
+    except Exception as e:
+        print(f"[QIE] GPU request skipped or failed: {e}")
+
     # 1) Ensure musubi-tuner is cloned before anything else
     _startup_clone_musubi_tuner()
 
@@ -331,4 +375,6 @@ if __name__ == "__main__":
 
     # 3) Launch Gradio app
     ui = build_ui()
+    # Limit concurrency (training is heavy). Enable queue for Spaces compatibility.
+    ui = ui.queue(concurrency_count=1, max_size=16)
     ui.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)))