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}")