Spaces:
Running on Zero
Running on Zero
| import io | |
| import json | |
| import os | |
| import re | |
| import atexit | |
| import threading | |
| import queue | |
| import time | |
| from datetime import datetime, timezone | |
| from urllib.parse import unquote | |
| import boto3 | |
| from botocore.exceptions import ClientError | |
| from PIL import Image | |
| FLUSH_INTERVAL_SEC = 15 | |
| FLUSH_MAX_ITEMS = 4 | |
| NUMBER_WIDTH = 4 | |
| INDEX_KEY = "_meta/index.json" | |
| FAVS_KEY = "_meta/favorites.json" | |
| _q = queue.Queue() | |
| _lock = threading.RLock() | |
| _worker_started = False | |
| _counters = {} | |
| _logs = {} | |
| def _safe(name): | |
| name = str(name or "unknown").strip() | |
| name = re.sub(r"[^\w.\-+ ]", "_", name) | |
| name = re.sub(r"\s+", "_", name) | |
| return name.strip("._ ")[:100] or "unknown" | |
| def _config_ready(): | |
| return all(os.environ.get(k) for k in ("R2_ENDPOINT", "R2_KEY", "R2_SECRET")) | |
| def is_enabled(): | |
| return _config_ready() | |
| def _client(): | |
| return boto3.client( | |
| "s3", | |
| endpoint_url=os.environ["R2_ENDPOINT"].rstrip("/"), | |
| aws_access_key_id=os.environ["R2_KEY"], | |
| aws_secret_access_key=os.environ["R2_SECRET"], | |
| region_name="auto", | |
| ) | |
| def _bucket(): | |
| return os.environ.get("R2_BUCKET", "imagegen") | |
| def _read_text(c, key): | |
| try: | |
| obj = c.get_object(Bucket=_bucket(), Key=key) | |
| return obj["Body"].read().decode("utf-8", errors="replace") | |
| except ClientError as e: | |
| if e.response.get("Error", {}).get("Code") in ("NoSuchKey", "404", "NotFound"): | |
| return "" | |
| raise | |
| def _read_index(c): | |
| raw = _read_text(c, INDEX_KEY) | |
| if not raw: | |
| return [] | |
| try: | |
| data = json.loads(raw) | |
| return data if isinstance(data, list) else [] | |
| except Exception: | |
| print("⚠️ [R2Storage] _meta/index.json is not valid JSON; keeping a fresh index.") | |
| return [] | |
| def _write_json(c, key, value): | |
| c.put_object( | |
| Bucket=_bucket(), Key=key, | |
| Body=json.dumps(value, ensure_ascii=False, separators=(",", ":")).encode("utf-8"), | |
| ContentType="application/json; charset=utf-8", | |
| ) | |
| def _init_folder(c, folder): | |
| if folder in _counters: | |
| return | |
| last = 0 | |
| log_text = _read_text(c, f"{folder}/log.txt") | |
| nums = re.findall(r"^#(\d+)\b", log_text, flags=re.M) | |
| if nums: | |
| last = max(int(n) for n in nums) | |
| try: | |
| paginator = c.get_paginator("list_objects_v2") | |
| for page in paginator.paginate(Bucket=_bucket(), Prefix=f"{folder}/"): | |
| for obj in page.get("Contents", []): | |
| match = re.match(rf"^{re.escape(folder)}/(\d+)\.png$", obj["Key"]) | |
| if match: | |
| last = max(last, int(match.group(1))) | |
| except Exception as e: | |
| print(f"⚠️ [R2Storage] Could not inspect folder {folder}: {e}") | |
| _counters[folder] = last | |
| _logs[folder] = log_text | |
| def _split_params(params): | |
| params = (params or "").strip() | |
| settings_text = "" | |
| idx = params.rfind("\nSteps: ") | |
| if idx != -1: | |
| settings_text = params[idx + 1:].strip() | |
| params = params[:idx] | |
| negative = "" | |
| nidx = params.find("\nNegative prompt: ") | |
| if nidx != -1: | |
| negative = params[nidx + len("\nNegative prompt: "):].strip() | |
| params = params[:nidx] | |
| return params.strip(), negative, settings_text | |
| def _parse_settings(s): | |
| out = {} | |
| # Common A1111-style values. This deliberately keeps parsing conservative. | |
| patterns = [ | |
| ("Steps", r"(?:^|,\s*)Steps:\s*([^,]+)"), | |
| ("Sampler", r"(?:^|,\s*)Sampler:\s*([^,]+)"), | |
| ("Scheduler", r"(?:^|,\s*)Scheduler:\s*([^,]+)"), | |
| ("CFG scale", r"(?:^|,\s*)CFG scale:\s*([^,]+)"), | |
| ("Seed", r"(?:^|,\s*)Seed:\s*([^,]+)"), | |
| ("Size", r"(?:^|,\s*)Size:\s*([^,]+)"), | |
| ("Base Model", r"(?:^|,\s*)Base Model:\s*(.*?)(?=,\s*[A-Z][^,:]*:|$)"), | |
| ] | |
| for key, pattern in patterns: | |
| m = re.search(pattern, s or "") | |
| if m: | |
| out[key] = m.group(1).strip() | |
| return out | |
| def _webp(data, max_side=None, quality=90): | |
| with Image.open(io.BytesIO(data)) as im: | |
| im = im.convert("RGB") | |
| if max_side: | |
| im.thumbnail((max_side, max_side)) | |
| buf = io.BytesIO() | |
| im.save(buf, "WEBP", quality=quality, method=4) | |
| return buf.getvalue(), im.size | |
| def _flush(items): | |
| if not items: | |
| return | |
| c = _client() | |
| touched = set() | |
| new_records = [] | |
| stamp = datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M UTC") | |
| with _lock: | |
| for it in items: | |
| folder = it["folder"] | |
| _init_folder(c, folder) | |
| _counters[folder] += 1 | |
| num = _counters[folder] | |
| filename = f"{num:0{NUMBER_WIDTH}d}.png" | |
| rel_path = f"{folder}/{filename}" | |
| base = rel_path[:-4] | |
| thumb_key = f"thumbs/{base}.webp" | |
| full_key = f"full/{base}.webp" | |
| prompt, negative, settings_text = _split_params(it.get("params", "")) | |
| settings = _parse_settings(settings_text) | |
| try: | |
| thumb, dimensions = _webp(it["data"], 512, 80) | |
| full, _ = _webp(it["data"], None, 92) | |
| c.put_object(Bucket=_bucket(), Key=rel_path, Body=it["data"], ContentType="image/png") | |
| c.put_object(Bucket=_bucket(), Key=thumb_key, Body=thumb, ContentType="image/webp") | |
| c.put_object(Bucket=_bucket(), Key=full_key, Body=full, ContentType="image/webp") | |
| entry = ( | |
| f"#{num:0{NUMBER_WIDTH}d} | {stamp} | {filename}\n" | |
| f"Prompt: {prompt}\n" | |
| f"Negative prompt: {negative}\n" | |
| f"Settings: {settings_text}\n" | |
| f"{'-' * 60}\n\n" | |
| ) | |
| _logs[folder] += entry | |
| touched.add(folder) | |
| model = it["model"] | |
| family = it["family"] | |
| size = settings.get("Size") or f"{dimensions[0]}x{dimensions[1]}" | |
| seed = settings.get("Seed", "") | |
| steps = settings.get("Steps", "") | |
| cfg = settings.get("CFG scale", "") | |
| try: steps = int(float(steps)) | |
| except (TypeError, ValueError): steps = 0 | |
| try: cfg = float(cfg) | |
| except (TypeError, ValueError): cfg = 0 | |
| orient = "مربع" if dimensions[0] == dimensions[1] else ("أفقي" if dimensions[0] > dimensions[1] else "عمودي") | |
| record = { | |
| "id": rel_path, | |
| "path": rel_path, | |
| "t": thumb_key, | |
| "f": full_key, | |
| "prompt": prompt, | |
| "negative": negative, | |
| "family": family, | |
| "model": model, | |
| "file": filename, | |
| "time": datetime.now(timezone.utc).isoformat(), | |
| "seed": seed, | |
| "size": size, | |
| "steps": steps, | |
| "sampler": settings.get("Sampler", ""), | |
| "scheduler": settings.get("Scheduler", ""), | |
| "cfg": cfg, | |
| "base_model": settings.get("Base Model", model), | |
| "orient": orient, | |
| "settings": settings, | |
| } | |
| new_records.append(record) | |
| except Exception as e: | |
| print(f"⚠️ [R2Storage] Failed uploading {rel_path}: {type(e).__name__}: {e}") | |
| for folder in touched: | |
| c.put_object( | |
| Bucket=_bucket(), Key=f"{folder}/log.txt", | |
| Body=_logs[folder].encode("utf-8"), | |
| ContentType="text/plain; charset=utf-8", | |
| ) | |
| if new_records: | |
| index = _read_index(c) | |
| existing_ids = {r.get("id") for r in index if isinstance(r, dict)} | |
| index.extend(r for r in new_records if r["id"] not in existing_ids) | |
| _write_json(c, INDEX_KEY, index) | |
| if new_records: | |
| print(f"☁️ [R2Storage] Saved {len(new_records)} image(s) directly to R2 bucket '{_bucket()}'.") | |
| def _worker(): | |
| pending = [] | |
| last_flush = time.time() | |
| while True: | |
| try: | |
| item = _q.get(timeout=2) | |
| if item is None: | |
| _flush(pending) | |
| return | |
| pending.append(item) | |
| except queue.Empty: | |
| pass | |
| due = time.time() - last_flush >= FLUSH_INTERVAL_SEC | |
| if pending and (due or len(pending) >= FLUSH_MAX_ITEMS): | |
| batch, pending = pending, [] | |
| try: | |
| _flush(batch) | |
| except Exception as e: | |
| print(f"⚠️ [R2Storage] Flush error: {type(e).__name__}: {e}") | |
| last_flush = time.time() | |
| def _shutdown(): | |
| if _worker_started: | |
| _q.put(None) | |
| time.sleep(1) | |
| def _ensure_worker(): | |
| global _worker_started | |
| if _worker_started: | |
| return | |
| with _lock: | |
| if _worker_started: | |
| return | |
| threading.Thread(target=_worker, daemon=True, name="r2-storage").start() | |
| atexit.register(_shutdown) | |
| _worker_started = True | |
| def queue_image(architecture, model, image_path, params): | |
| """Queue a generated image for direct upload to R2; never raises into generation.""" | |
| try: | |
| if not is_enabled(): | |
| print("ℹ️ [R2Storage] Skipped: R2_ENDPOINT/R2_KEY/R2_SECRET are not all configured.") | |
| return | |
| with open(image_path, "rb") as f: | |
| data = f.read() | |
| _ensure_worker() | |
| _q.put({ | |
| "folder": f"{_safe(architecture)}/{_safe(model)}", | |
| "family": _safe(architecture), | |
| "model": _safe(model), | |
| "data": data, | |
| "params": params or "", | |
| }) | |
| except Exception as e: | |
| print(f"⚠️ [R2Storage] Could not queue image: {type(e).__name__}: {e}") | |