| import os |
| import uuid |
| from io import BytesIO |
| from datetime import datetime, timezone, timedelta |
| from huggingface_hub import hf_hub_download |
|
|
| HF_TOKEN = os.environ.get("HF_LOGGING_TOKEN") |
| DATASET_REPO = os.environ.get("LOG_DATASET_REPO", "M3st3rJ4k3l/flux-klein-logs") |
| MAX_LOG_DAYS = int(os.environ.get("LOG_MAX_DAYS", "7")) |
|
|
|
|
| def _img_to_jpeg(img, quality=85): |
| if img is None: |
| return None |
| try: |
| buf = BytesIO() |
| img.convert("RGB").save(buf, format="JPEG", quality=quality) |
| return buf.getvalue() |
| except Exception: |
| return None |
|
|
|
|
| def _build_table(pil_inputs, output_pil, prompt, seed, steps, guidance_scale, |
| input_width, input_height, duration_seconds, success, error_message, |
| lora_titles, lora_weights, upscale_factor, lora_prompt_text, 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"}, |
| "lora_titles": {"feature": {"dtype": "string", "_type": "Value"}, "_type": "Sequence"}, |
| "lora_weights": {"feature": {"dtype": "float32", "_type": "Value"}, "_type": "Sequence"}, |
| "upscale_factor": {"dtype": "string", "_type": "Value"}, |
| "lora_prompt_text": {"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()), |
| ("lora_titles", pa.list_(pa.string())), |
| ("lora_weights", pa.list_(pa.float32())), |
| ("upscale_factor", pa.string()), |
| ("lora_prompt_text", 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()), |
| "lora_titles": pa.array([list(lora_titles or [])], type=pa.list_(pa.string())), |
| "lora_weights": pa.array([[float(w) for w in (lora_weights or [])]], type=pa.list_(pa.float32())), |
| "upscale_factor": pa.array([str(upscale_factor or "None")], type=pa.string()), |
| "lora_prompt_text": pa.array([str(lora_prompt_text or "")], type=pa.string()), |
| }, schema=schema) |
|
|
|
|
| def _upload_parquet(api, repo_id, table, path_in_repo): |
| import tempfile |
| import pyarrow.parquet as pq |
| tmp_path = None |
| try: |
| with tempfile.NamedTemporaryFile(suffix=".parquet", delete=False) as tmp: |
| tmp_path = tmp.name |
| pq.write_table(table, tmp_path) |
| print(f"[log] uploading {path_in_repo} ({os.path.getsize(tmp_path)//1024}KB)") |
| api.upload_file( |
| path_or_fileobj=tmp_path, path_in_repo=path_in_repo, |
| repo_id=repo_id, repo_type="dataset", |
| ) |
| print(f"[log] upload done — {repo_id}/{path_in_repo}") |
| finally: |
| if tmp_path: |
| try: |
| os.unlink(tmp_path) |
| except Exception as e: |
| print(f"[log] failed to delete temp file {tmp_path}: {e}") |
|
|
|
|
| def _make_path(now, uid): |
| return f"data/{now.strftime('%Y-%m-%d-%H%M%S')}-{uid}.parquet" |
|
|
|
|
| def _file_date(path): |
| return os.path.basename(path)[:10] |
|
|
|
|
| def _maybe_squash_history(api, repo_id, now): |
| marker = "metadata/last_squash.txt" |
| today = now.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") |
| print(f"[log] squashed history for {repo_id}") |
| api.upload_file( |
| path_or_fileobj=today.encode(), path_in_repo=marker, |
| repo_id=repo_id, repo_type="dataset", |
| ) |
| print(f"[log] updated squash marker: {today}") |
| except Exception as e: |
| print(f"[log] squash warning: {e}") |
|
|
|
|
| def _prune_old_files(api, repo_id, keep_days, now): |
| if keep_days <= 0: |
| return |
| cutoff = (now - timedelta(days=keep_days)).strftime("%Y-%m-%d") |
| try: |
| to_delete = [ |
| f.path |
| for f in api.list_repo_tree(repo_id, repo_type="dataset", path_in_repo="data") |
| if f.path.endswith(".parquet") and _file_date(f.path) < cutoff |
| ] |
| for path in to_delete: |
| api.delete_file(path_in_repo=path, repo_id=repo_id, repo_type="dataset") |
| print(f"[log] pruned: {path}") |
| if to_delete: |
| print(f"[log] pruned {len(to_delete)} old file(s)") |
| except Exception as e: |
| print(f"[log] prune warning: {e}") |
|
|
|
|
| def log_inference(pil_inputs, output_pil, prompt, seed, steps, guidance_scale, |
| input_width, input_height, duration_seconds, success, error_message="", |
| *, lora_titles=None, lora_weights=None, upscale_factor="None", |
| lora_prompt_text=""): |
| import time as _time |
| _t0 = _time.perf_counter() |
| if not HF_TOKEN or not DATASET_REPO: |
| print(f"[log] skipped — HF_LOGGING_TOKEN={'set' if HF_TOKEN else 'missing'}, " |
| f"LOG_DATASET_REPO={'set' if DATASET_REPO else 'missing'}") |
| return |
| try: |
| from huggingface_hub import HfApi |
| now = datetime.now(timezone.utc) |
|
|
| _t1 = _time.perf_counter() |
| table = _build_table(pil_inputs, output_pil, prompt, seed, steps, guidance_scale, |
| input_width, input_height, duration_seconds, success, error_message, |
| lora_titles, lora_weights, upscale_factor, lora_prompt_text, now) |
| print(f"[log] build_table: {_time.perf_counter() - _t1:.3f}s") |
|
|
| uid = uuid.uuid4().hex[:8] |
| path_in_repo = _make_path(now, uid) |
|
|
| _t2 = _time.perf_counter() |
| api = HfApi(token=HF_TOKEN) |
| api.create_repo(repo_id=DATASET_REPO, repo_type="dataset", private=True, exist_ok=True) |
| print(f"[log] create_repo: {_time.perf_counter() - _t2:.3f}s") |
|
|
| _t3 = _time.perf_counter() |
| _upload_parquet(api, DATASET_REPO, table, path_in_repo) |
| print(f"[log] upload_parquet: {_time.perf_counter() - _t3:.3f}s") |
|
|
| _t4 = _time.perf_counter() |
| _prune_old_files(api, DATASET_REPO, MAX_LOG_DAYS, now) |
| print(f"[log] prune_old_files: {_time.perf_counter() - _t4:.3f}s") |
|
|
| _t5 = _time.perf_counter() |
| _maybe_squash_history(api, DATASET_REPO, now) |
| print(f"[log] squash_history: {_time.perf_counter() - _t5:.3f}s") |
|
|
| except Exception as log_err: |
| import traceback as _tb |
| print(f"[log] WARNING: {log_err}\n{_tb.format_exc()}") |
| finally: |
| print(f"[log] log_inference total: {_time.perf_counter() - _t0:.3f}s") |
|
|
