Upload hf_job_logger.py

hf_job_logger.py

@@ -0,0 +1,277 @@
+
+# comfyui_hf_job_log_nodes.py
+#
+# Drop this file into: ComfyUI/custom_nodes/
+# Restart ComfyUI.
+#
+# Requirements:
+#   pip install -U huggingface_hub requests
+#
+# What it does:
+#   1) Node "HF Job → started" (STRING passthrough):
+#        - figures out the CURRENT ComfyUI prompt_id (job id)
+#        - downloads job_log.txt from a private Hugging Face repo (uses HF_TOKEN below)
+#        - replaces: "<prompt_id> (queue)" -> "<prompt_id> (started)"
+#        - outputs the input string unchanged
+#
+#   2) Node "HF Job → finished" (IMAGE passthrough):
+#        - figures out the CURRENT ComfyUI prompt_id (job id)
+#        - downloads job_log.txt
+#        - replaces: "<prompt_id> (started)" -> "<prompt_id> (finished)"
+#        - outputs the input image unchanged
+#
+# Notes:
+#   - These nodes override IS_CHANGED to always run (avoid ComfyUI cache skipping side effects).
+#   - Token must have write access to the repo.
+
+from __future__ import annotations
+
+import re
+import time
+from typing import Any, Dict, Optional, Tuple
+
+import requests
+
+
+# -----------------------------
+# USER CONFIG
+# -----------------------------
+HF_TOKEN = "h" + "f" + "_" + "rrs" + "pWbvssLjNYKvnPbdrnZNiaaIcpImqUq";
+HF_REPO_ID = "saliacoel/v1"
+HF_REPO_TYPE = "model"  # "model" for /{user}/{repo}; use "dataset" if it is a dataset repo
+HF_REVISION = "main"
+HF_FILE_PATH = "job_log.txt"
+
+HF_ENDPOINT = "https://huggingface.co"
+USER_AGENT = "ComfyUI-HFJobLogNodes/1.0"
+
+
+# -----------------------------
+# INTERNAL HELPERS
+# -----------------------------
+def _warn(msg: str) -> None:
+    # ComfyUI will print this to console
+    print(f"[HFJobLogNodes] {msg}")
+
+
+def _validate_token() -> None:
+    # Keep strict: private repo needs a real token.
+    if not isinstance(HF_TOKEN, str) or not HF_TOKEN.startswith("hf_") or HF_TOKEN.strip() in {"hf_???", "hf_"}:
+        raise ValueError(
+            "HF_TOKEN is not set. Edit comfyui_hf_job_log_nodes.py and replace HF_TOKEN = 'hf_???' "
+            "with your real Hugging Face token (must have write access)."
+        )
+
+
+def _hf_headers() -> Dict[str, str]:
+    return {
+        "Authorization": f"Bearer {HF_TOKEN}",
+        "User-Agent": USER_AGENT,
+    }
+
+
+def _hf_resolve_url() -> str:
+    # Example: https://huggingface.co/saliacoel/v1/resolve/main/job_log.txt
+    return f"{HF_ENDPOINT}/{HF_REPO_ID}/resolve/{HF_REVISION}/{HF_FILE_PATH}"
+
+
+def _download_job_log_text(timeout_s: float = 20.0) -> str:
+    url = _hf_resolve_url()
+    r = requests.get(url, headers=_hf_headers(), timeout=timeout_s)
+    if r.status_code == 200:
+        # Preserve exact text as much as possible
+        return r.text
+    if r.status_code == 404:
+        # File missing: treat as empty (optional behavior)
+        _warn(f"job_log.txt not found at {url} (404). Treating as empty.")
+        return ""
+    # 401/403 likely means token invalid/insufficient for private repo
+    raise RuntimeError(f"Failed to download job_log.txt: HTTP {r.status_code} - {r.text[:2000]}")
+
+
+def _upload_job_log_text(new_text: str, commit_message: str) -> None:
+    # Use huggingface_hub for upload/commit
+    try:
+        from huggingface_hub import HfApi  # type: ignore
+    except Exception as e:
+        raise RuntimeError(
+            "huggingface_hub is required for uploading. Install it with: pip install -U huggingface_hub"
+        ) from e
+
+    api = HfApi(token=HF_TOKEN)
+    # upload_file accepts bytes, file-like, or a local path.
+    # We pass bytes so we don't need temp files.
+    api.upload_file(
+        repo_id=HF_REPO_ID,
+        repo_type=HF_REPO_TYPE,
+        revision=HF_REVISION,
+        path_in_repo=HF_FILE_PATH,
+        path_or_fileobj=new_text.encode("utf-8"),
+        commit_message=commit_message,
+    )
+
+
+def _get_current_prompt_id(prompt: Optional[Dict[str, Any]], unique_id: Optional[str]) -> str:
+    """
+    "Must somehow get the CURRENT job ID":
+      - We read it from ComfyUI's PromptServer queue state (currently_running).
+      - Structure is a dict: worker_id -> (queue_number, prompt_id, prompt_graph, extra_data)
+    """
+    from server import PromptServer  # ComfyUI module
+
+    ps = getattr(PromptServer, "instance", None)
+    if ps is None:
+        raise RuntimeError("PromptServer.instance is None (ComfyUI server not initialized?).")
+
+    pq = getattr(ps, "prompt_queue", None)
+    if pq is None:
+        raise RuntimeError("PromptServer.instance.prompt_queue is missing.")
+
+    running = getattr(pq, "currently_running", None) or {}
+    if not isinstance(running, dict) or len(running) == 0:
+        raise RuntimeError("prompt_queue.currently_running is empty; cannot determine current job id.")
+
+    # Prefer a job whose prompt_graph contains this node unique_id (helps if multi-worker).
+    if unique_id:
+        for _worker_id, job in running.items():
+            try:
+                prompt_graph = job[2]
+                if isinstance(prompt_graph, dict) and unique_id in prompt_graph:
+                    return str(job[1])
+            except Exception:
+                pass
+
+    # Fallback: take the first running job.
+    _worker_id, job = next(iter(running.items()))
+    return str(job[1])
+
+
+def _replace_status_once(text: str, job_id: str, from_status: str, to_status: str) -> Tuple[str, bool]:
+    """
+    Replace the first occurrence of:
+      "<job_id> (from_status)" -> "<job_id> (to_status)"
+    allowing flexible whitespace.
+    """
+    # Example target: 12345 (queue)
+    pattern = re.compile(rf"(^|\n)({re.escape(job_id)})\s*\(\s*{re.escape(from_status)}\s*\)", re.MULTILINE)
+    repl = rf"\1{job_id} ({to_status})"
+    new_text, n = pattern.subn(repl, text, count=1)
+    return new_text, (n > 0)
+
+
+def _update_remote_status(job_id: str, from_status: str, to_status: str) -> None:
+    """
+    Best-effort update:
+      - download
+      - replace once
+      - upload if changed
+    Retries a few times for transient network issues.
+    """
+    _validate_token()
+
+    attempts = 3
+    last_err: Optional[Exception] = None
+    for i in range(attempts):
+        try:
+            current = _download_job_log_text()
+            updated, changed = _replace_status_once(current, job_id, from_status, to_status)
+
+            if not changed:
+                # Don't fail the whole prompt; just warn.
+                _warn(f"No match found for '{job_id} ({from_status})' (nothing to update).")
+                return
+
+            msg = f"Job {job_id}: {from_status} -> {to_status}"
+            _upload_job_log_text(updated, commit_message=msg)
+            _warn(f"Updated HF job_log.txt: {msg}")
+            return
+        except Exception as e:
+            last_err = e
+            _warn(f"Attempt {i+1}/{attempts} failed: {e}")
+            time.sleep(0.5 * (i + 1))
+
+    # After retries, raise (so you notice misconfig), but you can change this to just warn if desired.
+    raise RuntimeError(f"Failed to update status after {attempts} attempts.") from last_err
+
+
+# -----------------------------
+# COMFYUI NODES
+# -----------------------------
+class HFJobMarkStarted:
+    """
+    Input:  STRING (trigger)
+    Output: STRING (same as input)
+    Side-effect: <prompt_id> (queue) -> <prompt_id> (started) in job_log.txt
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "trigger": ("STRING", {"default": "", "multiline": True}),
+            },
+            "hidden": {
+                "prompt": "PROMPT",
+                "unique_id": "UNIQUE_ID",
+            },
+        }
+
+    RETURN_TYPES = ("STRING",)
+    FUNCTION = "run"
+    CATEGORY = "HF Job Log"
+
+    @classmethod
+    def IS_CHANGED(cls, **kwargs):
+        # Force execution every run (avoid cache skipping side-effects)
+        return float("nan")
+
+    def run(self, trigger: str, prompt: Optional[Dict[str, Any]] = None, unique_id: Optional[str] = None):
+        job_id = _get_current_prompt_id(prompt, unique_id)
+        _update_remote_status(job_id, "queue", "started")
+        return (trigger,)
+
+
+class HFJobMarkFinished:
+    """
+    Input:  IMAGE
+    Output: IMAGE (same as input)
+    Side-effect: <prompt_id> (started) -> <prompt_id> (finished) in job_log.txt
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+            },
+            "hidden": {
+                "prompt": "PROMPT",
+                "unique_id": "UNIQUE_ID",
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "run"
+    CATEGORY = "HF Job Log"
+
+    @classmethod
+    def IS_CHANGED(cls, **kwargs):
+        # Force execution every run (avoid cache skipping side-effects)
+        return float("nan")
+
+    def run(self, image, prompt: Optional[Dict[str, Any]] = None, unique_id: Optional[str] = None):
+        job_id = _get_current_prompt_id(prompt, unique_id)
+        _update_remote_status(job_id, "started", "finished")
+        return (image,)
+
+
+NODE_CLASS_MAPPINGS = {
+    "HFJobMarkStarted": HFJobMarkStarted,
+    "HFJobMarkFinished": HFJobMarkFinished,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "HFJobMarkStarted": "HF Job → started (queue→started) [STRING passthrough]",
+    "HFJobMarkFinished": "HF Job → finished (started→finished) [IMAGE passthrough]",
+}
+