Fix Space build and Gradio schema crash (state + deps)

app.py

@@ -16,6 +16,7 @@ from inference import _build_tokenizers, _resolve_device, load_model, run_infere
 RESULTS_DIR = "generated_results"
 DEFAULT_ANALYSIS_OUT = "analysis/outputs"
 os.makedirs(RESULTS_DIR, exist_ok=True)
+MODEL_CACHE = {}
 
 
 def discover_checkpoints():
@@ -131,6 +132,7 @@ def load_selected_model(checkpoint_label):
         "src_tok": src_tok,
         "tgt_tok": tgt_tok,
     }
+    MODEL_CACHE[checkpoint_label] = bundle
     model_info = {
         "checkpoint": ckpt_path,
         "experiment": experiment,
@@ -146,7 +148,19 @@ def load_selected_model(checkpoint_label):
     }
     status = f"Loaded `{experiment}` on `{device}` (`{cfg['model_type']}`)"
     suggested_out = os.path.join("analysis", "outputs_ui", experiment)
-    return bundle, status, json.dumps(model_info, ensure_ascii=False, indent=2), cfg["inference"]["num_steps"], suggested_out
+    return checkpoint_label, status, json.dumps(model_info, ensure_ascii=False, indent=2), cfg["inference"]["num_steps"], suggested_out
+
+
+def _get_bundle(model_key: str):
+    if not model_key:
+        raise gr.Error("Load a model first.")
+    if model_key not in MODEL_CACHE:
+        mapping = checkpoint_map()
+        if model_key not in mapping:
+            raise gr.Error("Selected checkpoint is no longer available. Refresh and load again.")
+        # Lazy reload if Space process restarted.
+        load_selected_model(model_key)
+    return MODEL_CACHE[model_key]
 
 
 def apply_preset(preset_name):
@@ -193,7 +207,7 @@ def save_generation(experiment, record):
 
 
 def generate_from_ui(
-    model_bundle,
+    model_key,
     input_text,
     temperature,
     top_k,
@@ -202,8 +216,7 @@ def generate_from_ui(
     num_steps,
     clean_output,
 ):
-    if not model_bundle:
-        raise gr.Error("Load a model first.")
+    model_bundle = _get_bundle(model_key)
     if not input_text.strip():
         raise gr.Error("Enter input text first.")
 
@@ -276,17 +289,15 @@ def _run_analysis_cmd(task, ckpt_path, output_dir, input_text="dharmo rakṣati
     return proc.returncode, log
 
 
-def run_single_task(model_bundle, task, output_dir, input_text, task4_phase):
-    if not model_bundle:
-        raise gr.Error("Load a model first.")
+def run_single_task(model_key, task, output_dir, input_text, task4_phase):
+    model_bundle = _get_bundle(model_key)
     code, log = _run_analysis_cmd(task, model_bundle["ckpt_path"], output_dir, input_text, task4_phase)
     status = f"Task {task} {'completed' if code == 0 else 'failed'} (exit={code})."
     return status, log
 
 
-def run_all_tasks(model_bundle, output_dir, input_text, task4_phase):
-    if not model_bundle:
-        raise gr.Error("Load a model first.")
+def run_all_tasks(model_key, output_dir, input_text, task4_phase):
+    model_bundle = _get_bundle(model_key)
     logs = []
     failures = 0
     for task in ["1", "2", "3", "4", "5"]:
@@ -349,7 +360,7 @@ CUSTOM_CSS = """
 
 
 with gr.Blocks(title="Sanskrit Diffusion Client Demo", css=CUSTOM_CSS) as demo:
-    model_state = gr.State(None)
+    model_state = gr.State("")
 
     gr.Markdown(
         """

requirements.txt

@@ -1,9 +1,7 @@
-gradio==5.23.3
-gradio_client==1.8.0
 torch>=2.2
 numpy>=1.24
 tqdm>=4.66
-huggingface_hub>=0.28,<1.0
+huggingface_hub==0.25.2
 tokenizers>=0.15
 datasets>=2.20
 scikit-learn>=1.4