Upload folder using huggingface_hub

app.py

@@ -38,6 +38,20 @@ except Exception:
     mlflow = None
 
 _MLFLOW_READY = False
+FLOW_STEPS = [
+    "Start",
+    "Load Model (checkpoint/config/device/eval)",
+    "Load Tokenizers",
+    "Input (IAST)",
+    "Source Tokenization",
+    "Encoder (run once)",
+    "KV Cache prepared",
+    "Initialize x_T (MASK)",
+    "Diffusion loop (T→0, with Task2/Task3 hooks)",
+    "Final x0",
+    "Decode to Devanagari",
+    "Evaluation/Tasks (Task4/Task5)",
+]
 
 
 def _setup_mlflow_once():
@@ -71,6 +85,41 @@ def _mlflow_event(run_name: str, params: dict | None = None, metrics: dict | Non
     except Exception:
         pass
 
+
+def _build_flow_markdown(model_loaded=False, inference_ready=False, task_states=None):
+    lines = ["### Execution Flow"]
+    for i, step in enumerate(FLOW_STEPS, start=1):
+        status = "⬜"
+        if model_loaded and i <= 3:
+            status = "✅"
+        if inference_ready and i <= 11:
+            status = "✅"
+        lines.append(f"{status} {i}. {step}")
+    if task_states:
+        lines.append("")
+        lines.append("### Task Status")
+        for k in ["1", "2", "3", "4", "5"]:
+            v = task_states.get(k, "pending")
+            icon = "✅" if v == "done" else ("🔄" if v.startswith("running") else ("❌" if v == "failed" else "⬜"))
+            lines.append(f"{icon} Task {k}: {v}")
+    return "\n".join(lines)
+
+
+def _task5_cfg(lambda_min, lambda_max, lambda_step, task5_samples):
+    lo = float(lambda_min)
+    hi = float(lambda_max)
+    st = max(0.1, float(lambda_step))
+    if hi < lo:
+        lo, hi = hi, lo
+    vals = []
+    cur = lo
+    while cur <= hi + 1e-9 and len(vals) < 30:
+        vals.append(round(cur, 2))
+        cur += st
+    if not vals:
+        vals = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0]
+    return {"scales": vals, "samples": max(5, int(task5_samples))}
+
 HF_DEFAULT_MODEL_REPO = os.environ.get("HF_DEFAULT_MODEL_REPO", "bhsinghgrid/DevaFlow")
 HF_DEFAULT_MODEL_FILE = os.environ.get("HF_DEFAULT_MODEL_FILE", "best_model.pt")
 
@@ -384,7 +433,7 @@ def _resolve_analysis_script() -> Path | None:
     return None
 
 
-def _run_analysis_cmd(task, ckpt_path, output_dir, input_text="dharmo rakṣati rakṣitaḥ", phase="analyze"):
+def _run_analysis_cmd(task, ckpt_path, output_dir, input_text="dharmo rakṣati rakṣitaḥ", phase="analyze", task5_samples=50):
     os.makedirs(output_dir, exist_ok=True)
     script = _resolve_analysis_script()
     if script is None:
@@ -420,6 +469,8 @@ def _run_analysis_cmd(task, ckpt_path, output_dir, input_text="dharmo rakṣati
         cmd.extend(["--input", input_text])
     if str(task) == "4":
         cmd.extend(["--phase", phase])
+    if str(task) == "5":
+        cmd.extend(["--task5_samples", str(int(task5_samples))])
 
     env = os.environ.copy()
     env.setdefault("HF_HOME", "/tmp/hf_home")
@@ -502,7 +553,7 @@ def _run_single_prediction(model_bundle, text: str, cfg_override: dict | None =
     return _decode_with_cleanup(tgt_tok, out[0].tolist(), text.strip(), cfg["inference"])
 
 
-def _live_task_analysis(model_bundle, task: str, input_text: str) -> str:
+def _live_task_analysis(model_bundle, task: str, input_text: str, task5_cfg: dict | None = None) -> str:
     text = input_text.strip()
     if not text:
         return "Live analysis skipped: empty input."
@@ -567,7 +618,7 @@ def _live_task_analysis(model_bundle, task: str, input_text: str) -> str:
 
     if str(task) == "5":
         ref = _iast_to_deva(text)
-        scales = [0.0, 0.5, 1.0, 1.5, 2.0]
+        scales = (task5_cfg or {}).get("scales", [0.0, 0.5, 1.0, 1.5, 2.0])
         rows = []
         for s in scales:
             cfg_map = {
@@ -582,13 +633,14 @@ def _live_task_analysis(model_bundle, task: str, input_text: str) -> str:
     return _live_input_summary(model_bundle, text)
 
 
-def _bg_worker(job_id: str, model_bundle, output_dir: str, input_text: str, task4_phase: str):
+def _bg_worker(job_id: str, model_bundle, output_dir: str, input_text: str, task4_phase: str, task5_cfg: dict):
     tasks = ["1", "2", "3", "4", "5"]
     failures = 0
     logs = []
     run_start = time.perf_counter()
     _BG_JOBS[job_id].update({"state": "running", "progress": 0, "failures": 0, "updated": datetime.now().isoformat()})
     for idx, task in enumerate(tasks, start=1):
+        _BG_JOBS[job_id]["task_states"][task] = "running"
         _BG_JOBS[job_id].update(
             {
                 "state": f"running task {task}",
@@ -596,13 +648,24 @@ def _bg_worker(job_id: str, model_bundle, output_dir: str, input_text: str, task
                 "updated": datetime.now().isoformat(),
             }
         )
-        code, log, used_bundled = _run_analysis_cmd(task, model_bundle["ckpt_path"], output_dir, input_text, task4_phase)
+        code, log, used_bundled = _run_analysis_cmd(
+            task,
+            model_bundle["ckpt_path"],
+            output_dir,
+            input_text,
+            task4_phase,
+            task5_cfg.get("samples", 50),
+        )
         logs.append(f"\n\n{'='*22} TASK {task} {'='*22}\n{log}")
         if code != 0:
             failures += 1
-            logs.append(f"\n[Live fallback]\n{_live_task_analysis(model_bundle, task, input_text)}\n")
+            _BG_JOBS[job_id]["task_states"][task] = "failed"
+            logs.append(f"\n[Live fallback]\n{_live_task_analysis(model_bundle, task, input_text, task5_cfg)}\n")
         elif used_bundled:
-            logs.append(f"\n[Live bundled summary]\n{_live_task_analysis(model_bundle, task, input_text)}\n")
+            _BG_JOBS[job_id]["task_states"][task] = "done(bundled)"
+            logs.append(f"\n[Live bundled summary]\n{_live_task_analysis(model_bundle, task, input_text, task5_cfg)}\n")
+        else:
+            _BG_JOBS[job_id]["task_states"][task] = "done"
         _BG_JOBS[job_id].update(
             {
                 "log": "".join(logs),
@@ -655,7 +718,7 @@ def _bg_worker(job_id: str, model_bundle, output_dir: str, input_text: str, task
     )
 
 
-def start_run_all_background(model_bundle, output_dir, input_text, task4_phase):
+def start_run_all_background(model_bundle, output_dir, input_text, task4_phase, task5_cfg):
     if not model_bundle:
         raise gr.Error("Load a model first.")
     os.makedirs(output_dir, exist_ok=True)
@@ -669,53 +732,64 @@ def start_run_all_background(model_bundle, output_dir, input_text, task4_phase):
         "output_dir": output_dir,
         "created": datetime.now().isoformat(),
         "updated": datetime.now().isoformat(),
+        "task_states": {k: "pending" for k in ["1", "2", "3", "4", "5"]},
     }
     th = threading.Thread(
         target=_bg_worker,
-        args=(job_id, model_bundle, output_dir, input_text, task4_phase),
+        args=(job_id, model_bundle, output_dir, input_text, task4_phase, task5_cfg),
         daemon=True,
     )
     th.start()
-    return f"Background run started. Job ID: {job_id}", f"Job {job_id} queued...", job_id
+    flow = _build_flow_markdown(model_loaded=True, inference_ready=False, task_states=_BG_JOBS[job_id]["task_states"])
+    return f"Background run started. Job ID: {job_id}", f"Job {job_id} queued...", job_id, _BG_JOBS[job_id]["task_states"], flow
 
 
 def poll_run_all_background(job_id, output_dir):
     if not job_id or job_id not in _BG_JOBS:
         msg = "Background job idle. You can run a single task or start Run All 5 in background."
         empty = refresh_task_outputs(output_dir)
-        return msg, msg, *empty
+        flow = _build_flow_markdown(model_loaded=False, inference_ready=False, task_states={})
+        return msg, msg, {}, flow, *empty
     j = _BG_JOBS[job_id]
     status = (
         f"Job {job_id} | state={j['state']} | progress={j['progress']}% | "
         f"failures={j['failures']} | updated={j['updated']}"
     )
     outputs = refresh_task_outputs(output_dir)
-    return status, j.get("log", ""), *outputs
+    flow = _build_flow_markdown(model_loaded=True, inference_ready=False, task_states=j.get("task_states", {}))
+    return status, j.get("log", ""), j.get("task_states", {}), flow, *outputs
 
 
-def run_single_task_and_refresh(model_bundle, task, output_dir, input_text, task4_phase):
-    status, log = run_single_task(model_bundle, task, output_dir, input_text, task4_phase)
+def run_single_task_and_refresh(model_bundle, task, output_dir, input_text, task4_phase, task5_cfg):
+    status, log, task_states, flow = run_single_task(model_bundle, task, output_dir, input_text, task4_phase, task5_cfg)
     out = refresh_task_outputs(output_dir)
-    return status, log, *out
+    return status, log, task_states, flow, *out
 
 
-def run_single_task(model_bundle, task, output_dir, input_text, task4_phase):
+def run_single_task(model_bundle, task, output_dir, input_text, task4_phase, task5_cfg):
     if not model_bundle:
         raise gr.Error("Load a model first.")
     t0 = time.perf_counter()
-    code, log, used_bundled = _run_analysis_cmd(task, model_bundle["ckpt_path"], output_dir, input_text, task4_phase)
+    code, log, used_bundled = _run_analysis_cmd(
+        task, model_bundle["ckpt_path"], output_dir, input_text, task4_phase, task5_cfg.get("samples", 50)
+    )
+    task_states = {k: "pending" for k in ["1", "2", "3", "4", "5"]}
+    task_states[str(task)] = "running"
     elapsed = (time.perf_counter() - t0) * 1000.0
     if code != 0:
         _bundle_task_outputs(model_bundle, output_dir)
-        log = f"{log}\n\n--- Live task analysis ---\n{_live_task_analysis(model_bundle, task, input_text)}"
+        log = f"{log}\n\n--- Live task analysis ---\n{_live_task_analysis(model_bundle, task, input_text, task5_cfg)}"
         status = f"Task {task} fallback mode: bundled reports + live input analysis."
+        task_states[str(task)] = "failed"
     else:
         if used_bundled:
             _bundle_task_outputs(model_bundle, output_dir)
-            log = f"{log}\n\n--- Live task analysis ---\n{_live_task_analysis(model_bundle, task, input_text)}"
+            log = f"{log}\n\n--- Live task analysis ---\n{_live_task_analysis(model_bundle, task, input_text, task5_cfg)}"
             status = f"Task {task} loaded from bundled analysis outputs + live analysis."
+            task_states[str(task)] = "done(bundled)"
         else:
             status = f"Task {task} completed (exit={code})."
+            task_states[str(task)] = "done"
     _mlflow_event(
         run_name=f"space_task_{task}",
         params={
@@ -731,10 +805,11 @@ def run_single_task(model_bundle, task, output_dir, input_text, task4_phase):
         },
         tags={"source": "hf_space", "mode": "single_task"},
     )
-    return status, log
+    flow = _build_flow_markdown(model_loaded=True, inference_ready=False, task_states=task_states)
+    return status, log, task_states, flow
 
 
-def run_all_tasks(model_bundle, output_dir, input_text, task4_phase):
+def run_all_tasks(model_bundle, output_dir, input_text, task4_phase, task5_cfg):
     if not model_bundle:
         raise gr.Error("Load a model first.")
     logs = []
@@ -742,7 +817,7 @@ def run_all_tasks(model_bundle, output_dir, input_text, task4_phase):
     used_bundled_any = False
     for task in ["1", "2", "3", "4", "5"]:
         code, log, used_bundled = _run_analysis_cmd(
-            task, model_bundle["ckpt_path"], output_dir, input_text, task4_phase
+            task, model_bundle["ckpt_path"], output_dir, input_text, task4_phase, task5_cfg.get("samples", 50)
         )
         logs.append(f"\n\n{'='*22} TASK {task} {'='*22}\n{log}")
         used_bundled_any = used_bundled_any or used_bundled
@@ -813,12 +888,16 @@ def _safe_refresh_task_outputs(output_dir):
         return (err, err, None, None, None, None, err, None, err, None)
 
 
-def _safe_start_run_all_background(model_bundle, output_dir, input_text, task4_phase, current_job_id):
+def _safe_start_run_all_background(
+    model_bundle, output_dir, input_text, task4_phase, current_job_id, lambda_min, lambda_max, lambda_step, task5_samples
+):
     try:
-        status, log, job_id = start_run_all_background(model_bundle, output_dir, input_text, task4_phase)
-        return status, log, job_id
+        cfg = _task5_cfg(lambda_min, lambda_max, lambda_step, task5_samples)
+        status, log, job_id, task_states, flow = start_run_all_background(model_bundle, output_dir, input_text, task4_phase, cfg)
+        return status, log, job_id, task_states, flow
     except Exception as e:
-        return f"Background start failed: {e}", f"Background start failed: {e}", current_job_id
+        err = f"Background start failed: {e}"
+        return err, err, current_job_id, {}, _build_flow_markdown(model_loaded=bool(model_bundle), inference_ready=False, task_states={})
 
 
 def _safe_poll_run_all_background(job_id, output_dir):
@@ -827,16 +906,43 @@ def _safe_poll_run_all_background(job_id, output_dir):
     except Exception as e:
         err = f"Track error: {e}"
         out = _safe_refresh_task_outputs(output_dir)
-        return err, err, *out
+        return err, err, {}, _build_flow_markdown(model_loaded=False, inference_ready=False, task_states={}), *out
 
 
-def _safe_run_single_task_and_refresh(model_bundle, task, output_dir, input_text, task4_phase):
+def _safe_run_single_task_and_refresh(
+    model_bundle, task, output_dir, input_text, task4_phase, lambda_min, lambda_max, lambda_step, task5_samples
+):
     try:
-        return run_single_task_and_refresh(model_bundle, task, output_dir, input_text, task4_phase)
+        cfg = _task5_cfg(lambda_min, lambda_max, lambda_step, task5_samples)
+        return run_single_task_and_refresh(model_bundle, task, output_dir, input_text, task4_phase, cfg)
     except Exception as e:
         err = f"Task {task} failed: {e}"
         out = _safe_refresh_task_outputs(output_dir)
-        return err, err, *out
+        return err, err, {}, _build_flow_markdown(model_loaded=bool(model_bundle), inference_ready=False, task_states={}), *out
+
+
+def _generate_with_flow(
+    model_bundle,
+    input_text,
+    temperature,
+    top_k,
+    repetition_penalty,
+    diversity_penalty,
+    num_steps,
+    clean_output,
+):
+    out_text, status, meta = generate_from_ui(
+        model_bundle,
+        input_text,
+        temperature,
+        top_k,
+        repetition_penalty,
+        diversity_penalty,
+        num_steps,
+        clean_output,
+    )
+    flow = _build_flow_markdown(model_loaded=True, inference_ready=True, task_states={})
+    return out_text, status, meta, flow
 
 
 CUSTOM_CSS = """
@@ -895,6 +1001,8 @@ with gr.Blocks(title="Sanskrit Diffusion Client Demo", css=CUSTOM_CSS) as demo:
     init_msg = "Select a model and load." if checkpoint_map() else "No checkpoints found in ablation_results/ or results*/."
     load_status = gr.Markdown(init_msg)
     model_info = gr.JSON(label="Loaded Model Details")
+    flow_box = gr.Markdown(_build_flow_markdown(model_loaded=False, inference_ready=False, task_states={}))
+    task_states_view = gr.JSON(label="Task Execution State (side-by-side)")
 
     with gr.Tabs():
         with gr.Tab("1) Task Runner"):
@@ -915,6 +1023,11 @@ with gr.Blocks(title="Sanskrit Diffusion Client Demo", css=CUSTOM_CSS) as demo:
                         value="analyze",
                         label="Task 4 Phase",
                     )
+                    gr.Markdown("**Task 5 Controls**")
+                    task5_lambda_min = gr.Slider(0.0, 3.0, value=0.0, step=0.1, label="Task5 λ min")
+                    task5_lambda_max = gr.Slider(0.0, 3.0, value=3.0, step=0.1, label="Task5 λ max")
+                    task5_lambda_step = gr.Slider(0.1, 1.0, value=0.5, step=0.1, label="Task5 λ step")
+                    task5_samples = gr.Slider(5, 200, value=50, step=5, label="Task5 sweep samples")
                     run_all_btn = gr.Button("Run All 5 Tasks (Background)", variant="primary")
                     track_bg_btn = gr.Button("Track Background Run")
 
@@ -1004,7 +1117,7 @@ with gr.Blocks(title="Sanskrit Diffusion Client Demo", css=CUSTOM_CSS) as demo:
     )
 
     generate_btn.click(
-        fn=generate_from_ui,
+        fn=_generate_with_flow,
         inputs=[
             model_state,
             input_text,
@@ -1015,10 +1128,10 @@ with gr.Blocks(title="Sanskrit Diffusion Client Demo", css=CUSTOM_CSS) as demo:
             num_steps,
             clean_output,
         ],
-        outputs=[output_text, run_status, run_record],
+        outputs=[output_text, run_status, run_record, flow_box],
     )
     input_text.submit(
-        fn=generate_from_ui,
+        fn=_generate_with_flow,
         inputs=[
             model_state,
             input_text,
@@ -1029,15 +1142,20 @@ with gr.Blocks(title="Sanskrit Diffusion Client Demo", css=CUSTOM_CSS) as demo:
             num_steps,
             clean_output,
         ],
-        outputs=[output_text, run_status, run_record],
+        outputs=[output_text, run_status, run_record, flow_box],
     )
 
     run_single_btn.click(
         fn=_safe_run_single_task_and_refresh,
-        inputs=[model_state, task_choice, analysis_output_dir, analysis_input, task4_phase],
+        inputs=[
+            model_state, task_choice, analysis_output_dir, analysis_input, task4_phase,
+            task5_lambda_min, task5_lambda_max, task5_lambda_step, task5_samples
+        ],
         outputs=[
             task_run_status,
             task_run_log,
+            task_states_view,
+            flow_box,
             task1_box,
             task2_box,
             task2_drift_img,
@@ -1052,8 +1170,11 @@ with gr.Blocks(title="Sanskrit Diffusion Client Demo", css=CUSTOM_CSS) as demo:
     )
     run_all_btn.click(
         fn=_safe_start_run_all_background,
-        inputs=[model_state, analysis_output_dir, analysis_input, task4_phase, bg_job_state],
-        outputs=[task_run_status, task_run_log, bg_job_state],
+        inputs=[
+            model_state, analysis_output_dir, analysis_input, task4_phase, bg_job_state,
+            task5_lambda_min, task5_lambda_max, task5_lambda_step, task5_samples
+        ],
+        outputs=[task_run_status, task_run_log, bg_job_state, task_states_view, flow_box],
     )
     track_bg_btn.click(
         fn=_safe_poll_run_all_background,
@@ -1061,6 +1182,8 @@ with gr.Blocks(title="Sanskrit Diffusion Client Demo", css=CUSTOM_CSS) as demo:
         outputs=[
             task_run_status,
             task_run_log,
+            task_states_view,
+            flow_box,
             task1_box,
             task2_box,
             task2_drift_img,