inference.py

#!/usr/bin/env python3
"""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 (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, os, sys, tempfile
import torch
from PIL import Image
from torchvision import transforms

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(HERE, "coz"))   # vendored Chain-of-Zoom modules

COZ_PROMPT = ("The second image is a zoom-in of the first image. Based on this knowledge, "
              "what is in the second image? Give me a set of words.")
_to_tensor = transforms.Compose([transforms.ToTensor()])


def resize_and_center_crop(img, size):
    w, h = img.size
    scale = size / min(w, h)
    nw, nh = int(w * scale), int(h * scale)
    img = img.resize((nw, nh), Image.LANCZOS)
    l, t = (nw - size) // 2, (nh - size) // 2
    return img.crop((l, t, l + size, t + size))


class _SRArgs:
    def __init__(self, ckpt, sd3, process_size):
        self.lora_path = f"{ckpt}/SR_LoRA/model_20001.pkl"
        self.vae_path = f"{ckpt}/SR_VAE/vae_encoder_20001.pt"
        self.pretrained_model_name_or_path = sd3
        self.process_size = process_size
        self.lora_rank = 4
        self.merge_and_unload_lora = False
        self.mixed_precision = "fp16"


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")
    sr.transformer.to("cuda:0", dtype=torch.float32)
    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)
    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):
    from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
    from qwen_vl_utils import process_vision_info
    from peft import PeftModel
    name = "Qwen/Qwen2.5-VL-3B-Instruct"
    model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
        name, torch_dtype="auto", device_map="auto", attn_implementation="sdpa")
    proc = AutoProcessor.from_pretrained(name)
    model = PeftModel.from_pretrained(model, f"{ckpt}/VLM_LoRA/checkpoint-10000").merge_and_unload().eval()
    return model, proc, process_vision_info


def vlm_prompt(model, proc, pvi, first, second, max_new_tokens=32):
    messages = [{"role": "system", "content": COZ_PROMPT},
                {"role": "user", "content": [{"type": "image", "image": first},
                                             {"type": "image", "image": second}]}]
    text = proc.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    ii, vi = pvi(messages)
    inputs = proc(text=[text], images=ii, videos=vi, padding=True, return_tensors="pt").to("cuda")
    gen = model.generate(**inputs, max_new_tokens=max_new_tokens)
    trimmed = [o[len(i):] for i, o in zip(inputs.input_ids, gen)]
    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(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")
    ap.add_argument("--rec_num", type=int, default=4)
    ap.add_argument("--upscale", type=int, default=4)
    ap.add_argument("--process_size", type=int, default=512)
    ap.add_argument("--max_new_tokens", type=int, default=32)
    a = ap.parse_args()

    sr = build_sr(a.ckpt, a.sd3, a.merged, a.process_size)
    model, proc, pvi = build_vlm(a.ckpt)
    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__":
    main()