Spaces:
Sleeping
Sleeping
| import os | |
| import uuid | |
| import threading | |
| import time as _time | |
| from io import BytesIO | |
| from datetime import datetime, timezone | |
| from huggingface_hub import hf_hub_download, CommitOperationAdd, CommitOperationDelete | |
| def _img_to_jpeg(img, quality=85): | |
| if img is None: | |
| return None | |
| buf = BytesIO() | |
| img.convert("RGB").save(buf, format="JPEG", quality=quality) | |
| return buf.getvalue() | |
| def _build_table(pil_inputs, output_pil, prompt, seed, steps, guidance_scale, | |
| input_width, input_height, duration_seconds, success, error_message, now): | |
| import json as _json | |
| import pyarrow as pa | |
| img_struct = pa.struct([("bytes", pa.binary()), ("path", pa.string())]) | |
| hf_meta = _json.dumps({"info": {"features": { | |
| "timestamp": {"dtype": "float64", "_type": "Value"}, | |
| "prompt": {"dtype": "string", "_type": "Value"}, | |
| "seed": {"dtype": "int32", "_type": "Value"}, | |
| "steps": {"dtype": "int32", "_type": "Value"}, | |
| "guidance_scale": {"dtype": "float32", "_type": "Value"}, | |
| "input_images": {"feature": {"_type": "Image"}, "_type": "Sequence"}, | |
| "output_image": {"_type": "Image"}, | |
| "duration_seconds": {"dtype": "float32", "_type": "Value"}, | |
| "input_width": {"dtype": "int32", "_type": "Value"}, | |
| "input_height": {"dtype": "int32", "_type": "Value"}, | |
| "success": {"dtype": "bool", "_type": "Value"}, | |
| "error_message": {"dtype": "string", "_type": "Value"}, | |
| }}}).encode() | |
| schema = pa.schema([ | |
| ("timestamp", pa.float64()), | |
| ("prompt", pa.string()), | |
| ("seed", pa.int32()), | |
| ("steps", pa.int32()), | |
| ("guidance_scale", pa.float32()), | |
| ("input_images", pa.list_(img_struct)), | |
| ("output_image", img_struct), | |
| ("duration_seconds", pa.float32()), | |
| ("input_width", pa.int32()), | |
| ("input_height", pa.int32()), | |
| ("success", pa.bool_()), | |
| ("error_message", pa.string()), | |
| ], metadata={b"huggingface": hf_meta}) | |
| def _img(b): | |
| return {"bytes": b, "path": None} | |
| input_jpegs = [_img_to_jpeg(img) for img in pil_inputs] | |
| output_jpeg = _img_to_jpeg(output_pil) | |
| return pa.table({ | |
| "timestamp": pa.array([now.timestamp()], type=pa.float64()), | |
| "prompt": pa.array([prompt], type=pa.string()), | |
| "seed": pa.array([int(seed)], type=pa.int32()), | |
| "steps": pa.array([int(steps)], type=pa.int32()), | |
| "guidance_scale": pa.array([float(guidance_scale)], type=pa.float32()), | |
| "input_images": pa.array([[_img(b) for b in input_jpegs]], type=pa.list_(img_struct)), | |
| "output_image": pa.array([_img(output_jpeg) if output_jpeg else None], type=img_struct), | |
| "duration_seconds": pa.array([float(duration_seconds)], type=pa.float32()), | |
| "input_width": pa.array([int(input_width)], type=pa.int32()), | |
| "input_height": pa.array([int(input_height)], type=pa.int32()), | |
| "success": pa.array([bool(success)], type=pa.bool_()), | |
| "error_message": pa.array([str(error_message)], type=pa.string()), | |
| }, schema=schema) | |
| def _write_parquet(table): | |
| import tempfile | |
| import pyarrow.parquet as pq | |
| with tempfile.NamedTemporaryFile(suffix=".parquet", delete=False) as tmp: | |
| path = tmp.name | |
| pq.write_table(table, path) | |
| return path | |
| def _make_path(now, uid): | |
| return f"data/{now.strftime('%Y-%m-%d-%H%M%S')}-{uid}.parquet" | |
| def _list_existing_files(api, repo_id): | |
| try: | |
| entries = list(api.list_repo_tree(repo_id, repo_type="dataset", path_in_repo="data")) | |
| except Exception as e: | |
| print(f"[log] could not list existing files (empty repo?): {e}") | |
| return [] | |
| return sorted(f.path for f in entries if f.path.endswith(".parquet")) | |
| def _build_add_ops(batch): | |
| return [CommitOperationAdd(path_in_repo=p, path_or_fileobj=local) | |
| for p, local in batch] | |
| def _build_delete_ops(existing_files, n_new, max_files): | |
| total_after = len(existing_files) + n_new | |
| if max_files <= 0 or total_after <= max_files: | |
| return [] | |
| n_delete = total_after - max_files | |
| return [CommitOperationDelete(path_in_repo=p) for p in existing_files[:n_delete]] | |
| def _delete_temp_files(batch): | |
| for _, local in batch: | |
| try: | |
| os.unlink(local) | |
| except Exception: | |
| pass | |
| def _squash_if_needed(api, repo_id): | |
| marker = "metadata/last_squash.txt" | |
| today = datetime.now(timezone.utc).strftime("%Y-%m-%d") | |
| try: | |
| try: | |
| local = hf_hub_download(repo_id=repo_id, filename=marker, | |
| repo_type="dataset", token=api.token) | |
| if open(local).read().strip() == today: | |
| return | |
| except Exception as e: | |
| print(f"[log] squash marker not found ({e}), proceeding with squash") | |
| api.super_squash_history(repo_id=repo_id, repo_type="dataset") | |
| api.upload_file(path_or_fileobj=today.encode(), path_in_repo=marker, | |
| repo_id=repo_id, repo_type="dataset") | |
| print(f"[log] squashed history for {repo_id}") | |
| except Exception as e: | |
| print(f"[log] squash warning: {e}") | |
| class LogUploader: | |
| def __init__(self, token, repo_id, max_files=5000, batch_interval=60): | |
| self._token = token | |
| self._repo_id = repo_id | |
| self._max_files = max_files | |
| self._batch_interval = batch_interval | |
| self._pending = [] | |
| self._lock = threading.Lock() | |
| if token and repo_id: | |
| threading.Thread(target=self._loop, daemon=True, name="log-uploader").start() | |
| def log_inference(self, pil_inputs, output_pil, prompt, seed, steps, guidance_scale, | |
| input_width, input_height, duration_seconds, success, error_message=""): | |
| if not self._token or not self._repo_id: | |
| print(f"[log] skipped — token={'set' if self._token else 'missing'}, repo={'set' if self._repo_id else 'missing'}") | |
| return | |
| t0 = _time.perf_counter() | |
| try: | |
| now = datetime.now(timezone.utc) | |
| table = _build_table(pil_inputs, output_pil, prompt, seed, steps, guidance_scale, | |
| input_width, input_height, duration_seconds, success, error_message, now) | |
| local_path = _write_parquet(table) | |
| path_in_repo = _make_path(now, uuid.uuid4().hex[:8]) | |
| self._enqueue(path_in_repo, local_path) | |
| print(f"[log] queued {path_in_repo} (pending={len(self._pending)})") | |
| except Exception as e: | |
| import traceback as _tb | |
| print(f"[log] WARNING: {e}\n{_tb.format_exc()}") | |
| print(f"[log] log_inference total: {_time.perf_counter() - t0:.3f}s") | |
| def _enqueue(self, path_in_repo, local_path): | |
| with self._lock: | |
| self._pending.append((path_in_repo, local_path)) | |
| def _drain(self): | |
| with self._lock: | |
| batch = self._pending[:] | |
| self._pending.clear() | |
| return batch | |
| def _requeue(self, batch): | |
| with self._lock: | |
| self._pending[:0] = batch | |
| def _loop(self): | |
| while True: | |
| _time.sleep(self._batch_interval) | |
| self._flush() | |
| def _flush(self): | |
| batch = self._drain() | |
| if not batch: | |
| return | |
| try: | |
| self._commit_batch(batch) | |
| _delete_temp_files(batch) | |
| except Exception as e: | |
| print(f"[log] batch upload warning: {e}") | |
| self._requeue(batch) | |
| def _commit_batch(self, batch): | |
| from huggingface_hub import HfApi | |
| api = HfApi(token=self._token) | |
| api.create_repo(repo_id=self._repo_id, repo_type="dataset", private=True, exist_ok=True) | |
| existing = _list_existing_files(api, self._repo_id) | |
| add_ops = _build_add_ops(batch) | |
| del_ops = _build_delete_ops(existing, len(batch), self._max_files) | |
| api.create_commit( | |
| repo_id=self._repo_id, repo_type="dataset", | |
| operations=add_ops + del_ops, | |
| commit_message=f"[log] batch {len(batch)}" + (f", prune {len(del_ops)}" if del_ops else ""), | |
| ) | |
| print(f"[log] committed {len(batch)} file(s), pruned {len(del_ops)}") | |
| _squash_if_needed(api, self._repo_id) | |