Upload inference.py with huggingface_hub

inference.py

@@ -1,14 +1,15 @@
 #!/usr/bin/env python3
-"""OracleZoom - faithful extreme super-resolution (4x -> 16x -> 64x -> 256x).
+"""OracleZoom - faithful extreme super-resolution of ONE image (4x -> 16x -> 64x -> 256x).
 
 Self-contained: this repo + auto-downloaded Stable Diffusion 3-medium and Qwen2.5-VL-3B.
 Vendored Chain-of-Zoom code lives in ./coz, checkpoints in ./ckpt, merged model = merged_transformer.safetensors.
 
-Usage:
-    python inference.py --input ./inputs --output ./outputs
-Outputs: outputs/per-scale/scale1..4/<name>.png  (4x / 16x / 64x / 256x)
+Usage (one image in, all scales out):
+    python inference.py --input photo.jpg --output ./outputs
+Writes: outputs/<name>_1x.png (input crop), _4x.png, _16x.png, _64x.png, _256x.png
+(To batch many images, just call zoom_image() in a loop.)
 """
-import argparse, glob, os, sys
+import argparse, os, sys, tempfile
 import torch
 from PIL import Image
 from torchvision import transforms
@@ -41,8 +42,9 @@ class _SRArgs:
         self.mixed_precision = "fp16"
 
 
-def build_sr(ckpt, sd3, process_size):
+def build_sr(ckpt, sd3, merged, process_size):
     from osediff_sd3 import OSEDiff_SD3_TEST, SD3Euler
+    from safetensors.torch import load_file
     sr = SD3Euler()
     for m in [sr.text_enc_1, sr.text_enc_2, sr.text_enc_3, sr.transformer, sr.vae]:
         m.to("cuda:0")
@@ -50,7 +52,14 @@ def build_sr(ckpt, sd3, process_size):
     sr.vae.to("cuda:0", dtype=torch.float32)
     for m in [sr.text_enc_1, sr.text_enc_2, sr.text_enc_3, sr.transformer, sr.vae]:
         m.requires_grad_(False)
-    return OSEDiff_SD3_TEST(_SRArgs(ckpt, sd3, process_size), sr)
+    sr_test = OSEDiff_SD3_TEST(_SRArgs(ckpt, sd3, process_size), sr)
+    # load the merged OracleZoom transformer (replaces all transformer weights)
+    sd = load_file(merged)
+    dev = next(sr_test.model.transformer.parameters()).device
+    sd = {k: v.to(dev, dtype=torch.float32) for k, v in sd.items()}
+    miss, unexp = sr_test.model.transformer.load_state_dict(sd, strict=False)
+    print(f"[OracleZoom] merged transformer loaded (missing={len(miss)} unexpected={len(unexp)})", flush=True)
+    return sr_test
 
 
 def build_vlm(ckpt):
@@ -77,10 +86,38 @@ def vlm_prompt(model, proc, pvi, first, second, max_new_tokens=32):
     return proc.batch_decode(trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
 
 
+def zoom_image(sr, model, proc, pvi, image_path, out_dir,
+               rec_num=4, upscale=4, process_size=512, max_new_tokens=32):
+    """Super-resolve ONE image through the recursion; save every scale to out_dir. Returns list of paths."""
+    os.makedirs(out_dir, exist_ok=True)
+    stem = os.path.splitext(os.path.basename(image_path))[0]
+    work = tempfile.mkdtemp()
+    resize_and_center_crop(Image.open(image_path).convert("RGB"), process_size).save(f"{work}/0.png")
+    for rec in range(rec_num):
+        prev = Image.open(f"{work}/{rec}.png").convert("RGB")
+        w, h = prev.size
+        nw, nh = w // upscale, h // upscale
+        crop = prev.crop(((w - nw) // 2, (h - nh) // 2, (w + nw) // 2, (h + nh) // 2))
+        zoom = crop.resize((w, h), Image.BICUBIC)
+        zoom.save(f"{work}/{rec + 1}_input.png")
+        prompt = vlm_prompt(model, proc, pvi, f"{work}/{rec}.png", f"{work}/{rec + 1}_input.png", max_new_tokens)
+        lq = _to_tensor(zoom).unsqueeze(0).to("cuda") * 2 - 1
+        with torch.no_grad():
+            out = torch.clamp(sr(lq, prompt=prompt)[0].cpu(), -1.0, 1.0)
+        transforms.ToPILImage()(out * 0.5 + 0.5).save(f"{work}/{rec + 1}.png")
+        print(f"  scale{rec + 1} ({4 ** (rec + 1)}x): {prompt}", flush=True)
+    saved = []
+    for s in range(rec_num + 1):
+        dst = os.path.join(out_dir, f"{stem}_{4 ** s}x.png")
+        Image.open(f"{work}/{s}.png").save(dst)
+        saved.append(dst)
+    return saved
+
+
 def main():
-    ap = argparse.ArgumentParser()
-    ap.add_argument("--input", required=True, help="folder of input images")
-    ap.add_argument("--output", required=True, help="output folder")
+    ap = argparse.ArgumentParser(description="OracleZoom: super-resolve one image to 4x/16x/64x/256x.")
+    ap.add_argument("--input", required=True, help="path to ONE input image")
+    ap.add_argument("--output", default="./outputs", help="output folder (all scales saved here)")
     ap.add_argument("--merged", default=os.path.join(HERE, "merged_transformer.safetensors"))
     ap.add_argument("--ckpt", default=os.path.join(HERE, "ckpt"))
     ap.add_argument("--sd3", default="stabilityai/stable-diffusion-3-medium-diffusers")
@@ -89,43 +126,12 @@ def main():
     ap.add_argument("--process_size", type=int, default=512)
     ap.add_argument("--max_new_tokens", type=int, default=32)
     a = ap.parse_args()
-    os.makedirs(a.output, exist_ok=True)
 
-    sr = build_sr(a.ckpt, a.sd3, a.process_size)
-    from safetensors.torch import load_file
-    sd = load_file(a.merged)
-    dev = next(sr.model.transformer.parameters()).device
-    sd = {k: v.to(dev, dtype=torch.float32) for k, v in sd.items()}
-    miss, unexp = sr.model.transformer.load_state_dict(sd, strict=False)
-    print(f"[OracleZoom] merged transformer loaded (missing={len(miss)} unexpected={len(unexp)})", flush=True)
+    sr = build_sr(a.ckpt, a.sd3, a.merged, a.process_size)
     model, proc, pvi = build_vlm(a.ckpt)
-
-    imgs = sorted(p for e in ("*.png", "*.jpg", "*.jpeg", "*.webp") for p in glob.glob(f"{a.input}/{e}"))
-    for img_path in imgs:
-        bname = os.path.splitext(os.path.basename(img_path))[0]
-        rec_dir = os.path.join(a.output, "per-sample", bname)
-        os.makedirs(rec_dir, exist_ok=True)
-        cur = resize_and_center_crop(Image.open(img_path).convert("RGB"), a.process_size)
-        cur.save(f"{rec_dir}/0.png")
-        for rec in range(a.rec_num):
-            prev = Image.open(f"{rec_dir}/{rec}.png").convert("RGB")
-            w, h = prev.size
-            nw, nh = w // a.upscale, h // a.upscale
-            crop = prev.crop(((w - nw) // 2, (h - nh) // 2, (w + nw) // 2, (h + nh) // 2))
-            zoom = crop.resize((w, h), Image.BICUBIC)
-            zp = f"{rec_dir}/{rec + 1}_input.png"
-            zoom.save(zp)
-            prompt = vlm_prompt(model, proc, pvi, f"{rec_dir}/{rec}.png", zp, a.max_new_tokens)
-            lq = _to_tensor(zoom).unsqueeze(0).to("cuda") * 2 - 1
-            with torch.no_grad():
-                out = torch.clamp(sr(lq, prompt=prompt)[0].cpu(), -1.0, 1.0)
-            transforms.ToPILImage()(out * 0.5 + 0.5).save(f"{rec_dir}/{rec + 1}.png")
-            print(f"  {bname}  scale{rec + 1} ({4 ** (rec + 1)}x): {prompt}", flush=True)
-        for s in range(a.rec_num + 1):
-            d = os.path.join(a.output, "per-scale", f"scale{s}")
-            os.makedirs(d, exist_ok=True)
-            Image.open(f"{rec_dir}/{s}.png").save(os.path.join(d, f"{bname}.png"))
-    print("[OracleZoom] done ->", a.output, flush=True)
+    saved = zoom_image(sr, model, proc, pvi, a.input, a.output,
+                       a.rec_num, a.upscale, a.process_size, a.max_new_tokens)
+    print("[OracleZoom] saved:", *saved, sep="\n  ", flush=True)
 
 
 if __name__ == "__main__":