Spaces:
Runtime error
Runtime error
File size: 9,432 Bytes
6c478b0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 | 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")
|