inference_coz_single.py

import os
import tempfile
import uuid
import torch
from PIL import Image
from torchvision import transforms
from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
from qwen_vl_utils import process_vision_info
from osediff_sd3 import OSEDiff_SD3_TEST, SD3Euler

# -------------------------------------------------------------------
# Helper: Resize & center-crop to a fixed square
# -------------------------------------------------------------------
def resize_and_center_crop(img: Image.Image, size: int) -> Image.Image:
    w, h = img.size
    scale = size / min(w, h)
    new_w, new_h = int(w * scale), int(h * scale)
    img = img.resize((new_w, new_h), Image.LANCZOS)
    left = (new_w - size) // 2
    top  = (new_h - size) // 2
    return img.crop((left, top, left + size, top + size))


# -------------------------------------------------------------------
# Helper: Generate a single VLM prompt for recursive_multiscale
# -------------------------------------------------------------------
def _generate_vlm_prompt(
    vlm_model,
    vlm_processor,
    process_vision_info,
    prev_image_path: str,
    zoomed_image_path: str,
    device: str = "cuda"
) -> str:
    """
    Given two image file paths:
      - prev_image_path:   the “full” image at the previous recursion.
      - zoomed_image_path: the cropped+resized (zoom) image for this step.
    This builds a single “recursive_multiscale” prompt via Qwen2.5-VL.
    Returns a string like “cat on sofa, pet, indoor, living room”, etc.
    """
    # (1) Define the system message for recursive_multiscale:
    message_text = (
        "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."
    )

    # (2) Build the two-image “chat” payload:
    messages = [
        {"role": "system", "content": message_text},
        {
            "role": "user",
            "content": [
                {"type": "image", "image": prev_image_path},
                {"type": "image", "image": zoomed_image_path},
            ],
        },
    ]

    # (3) Wrap through the VL processor to get “inputs”:
    text = vlm_processor.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    image_inputs, video_inputs = process_vision_info(messages)
    inputs = vlm_processor(
        text=[text],
        images=image_inputs,
        videos=video_inputs,
        padding=True,
        return_tensors="pt",
    ).to(device)

    # (4) Generate tokens → decode
    generated = vlm_model.generate(**inputs, max_new_tokens=128)
    # strip off the prompt tokens from each generated sequence:
    trimmed = [
        out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated)
    ]
    out_text = vlm_processor.batch_decode(
        trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
    )[0]

    # (5) Return exactly the bare words (no extra “,” if no additional user prompt)
    return out_text.strip()


# -------------------------------------------------------------------
# Main Function: recursive_multiscale_sr (with multiple centers)
# -------------------------------------------------------------------
def recursive_multiscale_sr(
    input_png_path: str,
    upscale: int,
    rec_num: int = 4,
    centers: list[tuple[float, float]] = None,
) -> tuple[list[Image.Image], list[str]]:
    """
    Perform `rec_num` recursive_multiscale super-resolution steps on a single PNG.
    - input_png_path: path to a single .png file on disk.
    - upscale: integer up-scale factor per recursion (e.g. 4).
    - rec_num: how many recursion steps to perform.
    - centers: a list of normalized (x, y) tuples in [0, 1], one per recursion step,
               indicating where to center the low-res crop for each step. The list
               length must equal rec_num. If centers is None, defaults to center=(0.5, 0.5) 
               for all steps.

    Returns a tuple (sr_pil_list, prompt_list), where:
      - sr_pil_list: list of PIL.Image outputs [SR1, SR2, …, SR_rec_num] in order.
      - prompt_list: list of the VLM prompts generated at each recursion.
    """
    ###############################
    # 0. Validate / fill default centers
    ###############################
    if centers is None:
        # Default: use center (0.5, 0.5) for every recursion
        centers = [(0.5, 0.5) for _ in range(rec_num)]
    else:
        if not isinstance(centers, (list, tuple)) or len(centers) != rec_num:
            raise ValueError(
                f"`centers` must be a list of {rec_num} (x,y) tuples, but got length {len(centers)}."
            )

    ###############################
    # 1. Fixed hyper-parameters
    ###############################
    device = "cuda"
    process_size = 512     # same as args.process_size

    # model checkpoint paths (hard-coded to your example)
    LORA_PATH = "ckpt/SR_LoRA/model_20001.pkl"
    VAE_PATH  = "ckpt/SR_VAE/vae_encoder_20001.pt"
    SD3_MODEL = "stabilityai/stable-diffusion-3-medium-diffusers"
    # VLM model name (hard-coded)
    VLM_NAME  = "Qwen/Qwen2.5-VL-3B-Instruct"

    ###############################
    # 2. Build a dummy “args” namespace
    #    to satisfy OSEDiff_SD3_TEST constructor.
    ###############################
    class _Args:
        pass

    args = _Args()
    args.upscale                       = upscale
    args.lora_path                     = LORA_PATH
    args.vae_path                      = VAE_PATH
    args.pretrained_model_name_or_path = SD3_MODEL
    args.merge_and_unload_lora         = False
    args.lora_rank                     = 4
    args.vae_decoder_tiled_size        = 224
    args.vae_encoder_tiled_size        = 1024
    args.latent_tiled_size             = 96
    args.latent_tiled_overlap          = 32
    args.mixed_precision               = "fp16"
    args.efficient_memory              = False
    # (other flags are not used by OSEDiff_SD3_TEST, so we skip them)

    ###############################
    # 3. Load the SD3 SR model (non-efficient)
    ###############################
    # 3.1 Instantiate the underlying SD3-Euler UNet/VAE/text encoders
    sd3 = SD3Euler()
    # move all text encoders + transformer + VAE to CUDA:
    sd3.text_enc_1.to(device)
    sd3.text_enc_2.to(device)
    sd3.text_enc_3.to(device)
    sd3.transformer.to(device, dtype=torch.float32)
    sd3.vae.to(device, dtype=torch.float32)
    # freeze
    for p in (
        sd3.text_enc_1,
        sd3.text_enc_2,
        sd3.text_enc_3,
        sd3.transformer,
        sd3.vae,
    ):
        p.requires_grad_(False)

    # 3.2 Wrap in OSEDiff_SD3_TEST helper:
    model_test = OSEDiff_SD3_TEST(args, sd3)
    # (by default, “model_test(...)” takes (lq_tensor, prompt=str) and returns a list[tensor])

    ###############################
    # 4. Load the VLM (Qwen2.5-VL)
    ###############################
    vlm_model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
        VLM_NAME,
        torch_dtype="auto",
        device_map="auto"   # immediately dispatches layers onto available GPUs
    )
    vlm_processor = AutoProcessor.from_pretrained(VLM_NAME)

    ###############################
    # 5. Pre-allocate a Temporary Directory
    #    to hold intermediate JPEG/PNG files
    ###############################
    unique_id = uuid.uuid4().hex
    prefix = f"recms_{unique_id}_"

    with tempfile.TemporaryDirectory(prefix=prefix) as td:
        # (we’ll write “prev.png” and “zoom.png” at each step)

        ###############################
        # 6. Prepare the very first “full” image
        ###############################
        # 6.1 Load + center crop → first_image is (512×512) PIL on CPU
        img0 = Image.open(input_png_path).convert("RGB")
        img0 = resize_and_center_crop(img0, process_size)

        # 6.2 Save it once so VLM can read it as “prev.png”
        prev_path = os.path.join(td, "step0_prev.png")
        img0.save(prev_path)

        # We will maintain lists of PIL outputs and prompts:
        sr_pil_list: list[Image.Image] = []
        prompt_list: list[str] = []

        ###############################
        # 7. Recursion loop (now up to rec_num times)
        ###############################
        for rec in range(rec_num):
            # (A) Load the previous SR output (or original) and compute crop window
            prev_pil = Image.open(prev_path).convert("RGB")
            w, h = prev_pil.size  # should be (512×512) each time

            # (1) Compute the “low-res” window size:
            new_w, new_h = w // upscale, h // upscale  # e.g. 128×128 for upscale=4

            # (2) Map normalized center → pixel center, then clamp so crop stays in bounds:
            cx_norm, cy_norm = centers[rec]
            cx = int(cx_norm * w)
            cy = int(cy_norm * h)
            half_w = new_w // 2
            half_h = new_h // 2

            # If center in pixels is too close to left/top, clamp so left=0 or top=0; same on right/bottom
            left = cx - half_w
            top = cy - half_h
            # clamp left ∈ [0, w - new_w], top ∈ [0, h - new_h]
            left = max(0, min(left, w - new_w))
            top = max(0, min(top, h - new_h))
            right = left + new_w
            bottom = top + new_h

            cropped = prev_pil.crop((left, top, right, bottom))

            # (B) Resize that crop back up to (512×512) via BICUBIC → zoomed
            zoomed = cropped.resize((w, h), Image.BICUBIC)
            zoom_path = os.path.join(td, f"step{rec+1}_zoom.png")
            zoomed.save(zoom_path)

            # (C) Generate a recursive_multiscale VLM “tag” prompt
            prompt_tag = _generate_vlm_prompt(
                vlm_model=vlm_model,
                vlm_processor=vlm_processor,
                process_vision_info=process_vision_info,
                prev_image_path=prev_path,
                zoomed_image_path=zoom_path,
                device=device,
            )
            # (By default, no extra user prompt is appended.)

            # (D) Prepare the low-res tensor for SR: convert zoomed → Tensor → [0,1] → [−1,1]
            to_tensor = transforms.ToTensor()
            lq = to_tensor(zoomed).unsqueeze(0).to(device)  # shape (1,3,512,512)
            lq = (lq * 2.0) - 1.0

            # (E) Do SR inference:
            with torch.no_grad():
                out_tensor = model_test(lq, prompt=prompt_tag)[0]  # (3,512,512) on CPU or GPU
                out_tensor = out_tensor.clamp(-1.0, 1.0).cpu()
                # back to PIL in [0,1]:
                out_pil = transforms.ToPILImage()((out_tensor * 0.5) + 0.5)

            # (F) Save this step’s SR output as “prev.png” for next iteration:
            out_path = os.path.join(td, f"step{rec+1}_sr.png")
            out_pil.save(out_path)
            prev_path = out_path

            # (G) Append the PIL to our list:
            sr_pil_list.append(out_pil)
            prompt_list.append(prompt_tag)

        # end for(rec)

        ###############################
        # 8. Return the SR outputs & prompts
        ###############################
        # The list sr_pil_list = [ SR1, SR2, …, SR_rec_num ] in order.
        return sr_pil_list, prompt_list