VLM lora added

inference_coz_single.py

@@ -7,6 +7,7 @@ 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
+from peft import PeftModel
 
 # -------------------------------------------------------------------
 # Helper: Resize & center-crop to a fixed square
@@ -95,6 +96,60 @@ def _generate_vlm_prompt(
     return out_text.strip()
 
 
+VLM_NAME  = "Qwen/Qwen2.5-VL-3B-Instruct"
+
+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)
+
+vlm_model = PeftModel.from_pretrained(vlm_model, "ckpt/VLM_LoRA/checkpoint-10000")
+vlm_model = vlm_model.merge_and_unload()
+vlm_model.eval()
+
+device = "cuda"
+process_size = 512
+
+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"
+
+class _Args:
+    pass
+
+args = _Args()
+args.upscale                       = 4
+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
+
+sd3 = SD3Euler()
+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)
+
+for p in (
+    sd3.text_enc_1,
+    sd3.text_enc_2,
+    sd3.text_enc_3,
+    sd3.transformer,
+    sd3.vae,
+):
+    p.requires_grad_(False)
+
+model_test = OSEDiff_SD3_TEST(args, sd3)
 
 # -------------------------------------------------------------------
 # Main Function: recursive_multiscale_sr (with multiple centers)
@@ -131,101 +186,23 @@ def recursive_multiscale_sr(
                 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 (512×512)
         img0 = Image.open(input_png_path).convert("RGB")
         img0 = resize_and_center_crop(img0, process_size)
 
-        # Note: we no longer need to write “prev.png” to disk. Just keep it in memory.
+
         prev_pil = img0.copy()
 
         sr_pil_list: list[Image.Image] = []
         prompt_list:  list[str]        = []
 
         for rec in range(rec_num):
-            # (A) Compute low-res crop window on prev_pil
+
             w, h = prev_pil.size  # (512×512)
             new_w, new_h = w // upscale, h // upscale
 
@@ -240,10 +217,10 @@ def recursive_multiscale_sr(
 
             cropped = prev_pil.crop((left, top, right, bottom))
 
-            # (B) Upsample that crop back to (512×512)
+
             zoomed_pil = cropped.resize((w, h), Image.BICUBIC)
 
-            # (C) Generate VLM prompt by passing PILs directly:
+
             prompt_tag = _generate_vlm_prompt(
                 vlm_model=vlm_model,
                 vlm_processor=vlm_processor,
@@ -253,22 +230,22 @@ def recursive_multiscale_sr(
                 device=device,
             )
 
-            # (D) Prepare “zoomed_pil” → tensor in [−1, 1]
+
             to_tensor = transforms.ToTensor()
             lq = to_tensor(zoomed_pil).unsqueeze(0).to(device)  # (1,3,512,512)
             lq = (lq * 2.0) - 1.0
 
-            # (E) Run SR inference
+
             with torch.no_grad():
                 out_tensor = model_test(lq, prompt=prompt_tag)[0]
                 out_tensor = out_tensor.clamp(-1.0, 1.0).cpu()
                 out_pil = transforms.ToPILImage()((out_tensor * 0.5) + 0.5)
 
-            # (F) Bookkeeping: set prev_pil = out_pil for next iteration
+
             prev_pil = out_pil
 
             # (G) Append to results
             sr_pil_list.append(out_pil)
             prompt_list.append(prompt_tag)
 
-        return sr_pil_list, prompt_list
+        return sr_pil_list, prompt_list