Archive images to fallback bucket

app.py

@@ -44,7 +44,11 @@ from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import FileResponse, Response
 from gradio_client import Client
 from gradio.routes import App
-from huggingface_hub import download_bucket_files, list_bucket_tree
+from huggingface_hub import (
+    batch_bucket_files,
+    download_bucket_files,
+    list_bucket_tree,
+)
 from PIL import Image, ImageDraw, ImageFont
 from PIL.PngImagePlugin import PngInfo
 from pydantic import BaseModel, ConfigDict, Field
@@ -74,7 +78,7 @@ os.environ.setdefault("YOLO_CONFIG_DIR", "/tmp/Ultralytics")
 
 DATA_MOUNT = Path("/data")
 BUCKET_MOUNT = Path("/CB")
-HAS_DATA_MOUNT = DATA_MOUNT.is_dir()
+HAS_DATA_MOUNT = (DATA_MOUNT / "img").is_dir()
 DEFAULT_DATA_DIR = DATA_MOUNT if HAS_DATA_MOUNT else Path.cwd() / "data"
 DATA_DIR = Path(os.environ.get("DATA_DIR", DEFAULT_DATA_DIR))
 LOCAL_MODEL_DIR = Path(os.environ.get("LOCAL_MODEL_DIR", "/tmp/models"))
@@ -139,6 +143,9 @@ if not PASSWORD:
     raise RuntimeError("pass environment variable is required")
 
 BUCKET_ID = "HyperHail/CB"
+IMAGE_BUCKET_ID = "HyperHail/C"
+IMAGE_BUCKET_PREFIX = "img"
+IMAGE_TOKEN = os.environ.get("hh") or False
 MODEL_KINDS = (
     "checkpoints", "diffusion_models", "clip", "vae", "loras",
     "upscale_models", "ultralytics",
@@ -263,6 +270,11 @@ def cleanup_mount():
             if path.is_file() and path.suffix.casefold() == ".7z":
                 path.unlink()
                 log(f"Deleted img/{path.relative_to(IMAGE_DIR)}")
+        if HAS_DATA_MOUNT:
+            for path in IMAGE_DIR.glob("????-??-??/*.png"):
+                with Image.open(path) as image:
+                    save_named_image(image, path.with_suffix(".epng"))
+                path.unlink()
 
 
 def download(url, target):
@@ -1538,10 +1550,11 @@ def generate_image(request, from_api=True):
 
 
 def archive_image(image, from_api):
-    if not HAS_DATA_MOUNT:
-        return
     try:
-        save_image(image, from_api)
+        if HAS_DATA_MOUNT:
+            save_image(image, from_api)
+        else:
+            upload_image(image)
     except Exception as error:
         log(f"Archive failed: {error}")
 
@@ -1566,29 +1579,46 @@ def image_key(salt):
 def save_image(image, from_api):
     path = IMAGE_DIR / datetime.now(TIMEZONE).date().isoformat()
     path.mkdir(parents=True, exist_ok=True)
-    suffix = "ST.png" if from_api else "C.epng"
-    number = (
-        sum(
-            file.name.endswith(suffix)
-            and file.name.removesuffix(suffix).isdigit()
+    suffix = "ST.epng" if from_api else "C.epng"
+    number = max(
+        (
+            int(file.name.removesuffix(suffix))
             for file in path.iterdir()
-        )
-        + 1
+            if file.name.endswith(suffix)
+            and file.name.removesuffix(suffix).isdigit()
+        ),
+        default=0,
+    ) + 1
+    save_named_image(image, path / f"{number}{suffix}")
+
+
+def upload_image(image):
+    date = datetime.now(TIMEZONE).date().isoformat()
+    folder = f"{IMAGE_BUCKET_PREFIX}/{date}"
+    suffix = "ST.epng"
+    number = max(
+        (
+            int(Path(item.path).name.removesuffix(suffix))
+            for item in list_bucket_tree(
+                IMAGE_BUCKET_ID,
+                prefix=f"{folder}/",
+                recursive=True,
+                token=IMAGE_TOKEN,
+            )
+            if item.type == "file"
+            and Path(item.path).name.endswith(suffix)
+            and Path(item.path).name.removesuffix(suffix).isdigit()
+        ),
+        default=0,
+    ) + 1
+    batch_bucket_files(
+        IMAGE_BUCKET_ID,
+        add=[(encrypted_image_bytes(image), f"{folder}/{number}{suffix}")],
+        token=IMAGE_TOKEN,
     )
-    if from_api:
-        name = f"{number}{suffix}"
-        output = path / name
-        temp = output.with_suffix(".png.part")
-        temp.write_bytes(image_png_bytes(image))
-        temp.replace(output)
-        log(f"Saved img/{output.relative_to(IMAGE_DIR)}")
-    else:
-        save_named_image(image, path / f"{number}{suffix}")
 
 
-def save_named_image(image, output):
-    output.parent.mkdir(parents=True, exist_ok=True)
-    temp = output.with_suffix(".epng.part")
+def encrypted_image_bytes(image):
     salt = os.urandom(SALT_SIZE)
     nonce = os.urandom(NONCE_SIZE)
     encrypted = AESGCM(image_key(salt)).encrypt(
@@ -1596,7 +1626,13 @@ def save_named_image(image, output):
         image_png_bytes(image),
         FILE_MAGIC,
     )
-    temp.write_bytes(FILE_MAGIC + salt + nonce + encrypted)
+    return FILE_MAGIC + salt + nonce + encrypted
+
+
+def save_named_image(image, output):
+    output.parent.mkdir(parents=True, exist_ok=True)
+    temp = output.with_suffix(".epng.part")
+    temp.write_bytes(encrypted_image_bytes(image))
     temp.replace(output)
     log(f"Saved img/{output.relative_to(IMAGE_DIR)}")
 
@@ -2620,9 +2656,8 @@ def archive(
     body: bytes = Body(),
     from_api: bool = Query(True),
 ):
-    if HAS_DATA_MOUNT:
-        image = Image.open(BytesIO(body)).copy()
-        backup_pool.submit(save_image, image, from_api).result()
+    image = Image.open(BytesIO(body)).copy()
+    backup_pool.submit(archive_image, image, from_api).result()
     return {"status": True}