elephantmipt/trackio_stable-diffusion-v1-5_stable-diffusion-v1-5_None

1import threading2import time3from collections import deque4 5import huggingface_hub6from gradio_client import Client7 8from trackio.utils import RESERVED_KEYS, fibo, generate_readable_name9 10 11class Run:12    def __init__(13        self,14        url: str,15        project: str,16        client: Client,17        name: str | None = None,18        config: dict | None = None,19    ):20        self.url = url21        self.project = project22        self._client_lock = threading.Lock()23        self._client_thread = None24        self._client = client25        self.name = name or generate_readable_name()26        self.config = config or {}27        self._queued_logs = deque()28 29        if client is None:30            self._client_thread = threading.Thread(target=self._init_client_background)31            self._client_thread.start()32 33    def _init_client_background(self):34        fib = fibo()35        for sleep_coefficient in fib:36            try:37                client = Client(self.url, verbose=False)38                with self._client_lock:39                    self._client = client40                    if len(self._queued_logs) > 0:41                        for queued_log in self._queued_logs:42                            self._client.predict(**queued_log)43                        self._queued_logs.clear()44                    break45            except Exception:46                pass47            if sleep_coefficient is not None:48                time.sleep(0.1 * sleep_coefficient)49 50    def log(self, metrics: dict):51        for k in metrics.keys():52            if k in RESERVED_KEYS or k.startswith("__"):53                raise ValueError(54                    f"Please do not use this reserved key as a metric: {k}"55                )56        with self._client_lock:57            if self._client is None:58                # client can still be None for a Space while the Space is still initializing.59                # queue up log items for when the client is not None.60                self._queued_logs.append(61                    dict(62                        api_name="/log",63                        project=self.project,64                        run=self.name,65                        metrics=metrics,66                        hf_token=huggingface_hub.utils.get_token(),67                    )68                )69            else:70                assert (71                    len(self._queued_logs) == 072                )  # queue should have been flushed on client init73                # write the current log item74                self._client.predict(75                    api_name="/log",76                    project=self.project,77                    run=self.name,78                    metrics=metrics,79                    hf_token=huggingface_hub.utils.get_token(),80                )81 82    def finish(self):83        """Cleanup when run is finished."""84        # wait for background client thread, in case it has a queue of logs to flush.85        if self._client_thread is not None:86            print(f"* Uploading logs to Trackio Space: {self.url} (please wait...)")87            self._client_thread.join()88