import os import shutil import subprocess from pathlib import Path from PIL import Image import gradio as gr import spaces INPUT_DIR = "samples" OUTPUT_DIR = "inference_results/coz_vlmprompt" 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)) def make_preview_with_boxes(image_path: str, scale_option: str) -> Image.Image: try: orig = Image.open(image_path).convert("RGB") except Exception as e: fallback = Image.new("RGB", (512, 512), (200, 200, 200)) from PIL import ImageDraw draw = ImageDraw.Draw(fallback) draw.text((20, 20), f"Error:\n{e}", fill="red") return fallback base = resize_and_center_crop(orig, 512) scale_int = int(scale_option.replace("x", "")) if scale_int == 1: sizes = [512] * 4 else: sizes = [512 // (scale_int * (2 ** i)) for i in range(4)] from PIL import ImageDraw draw = ImageDraw.Draw(base) colors = ["red", "lime", "cyan", "yellow"] width = 3 for idx, s in enumerate(sizes): x0 = (512 - s) // 2 y0 = (512 - s) // 2 x1 = x0 + s y1 = y0 + s draw.rectangle([(x0, y0), (x1, y1)], outline=colors[idx], width=width) return base @spaces.GPU(duration=120) def run_with_upload(uploaded_image_path, upscale_option, session_id=None): """ Each invocation creates/uses: - samples//input.png ← user’s uploaded image - inference_results/coz_vlmprompt//per-sample/input/*.png ← inference outputs """ if uploaded_image_path is None: return [] # 1) Prepare a per-session input directory print(session_id) session_folder = os.path.join(INPUT_DIR, str(session_id)) os.makedirs(session_folder, exist_ok=True) # 2) Clear only this session’s folder for fn in os.listdir(session_folder): full_path = os.path.join(session_folder, fn) if os.path.isfile(full_path) or os.path.islink(full_path): os.remove(full_path) elif os.path.isdir(full_path): shutil.rmtree(full_path) # 3) Save uploaded image to session_folder/input.png try: pil_img = Image.open(uploaded_image_path).convert("RGB") save_path = Path(session_folder) / "input.png" pil_img.save(save_path, format="PNG") except Exception as e: print(f"Error: could not save uploaded image: {e}") return [] # 4) Define a per-session output directory session_output_dir = os.path.join(OUTPUT_DIR, str(session_id)) os.makedirs(session_output_dir, exist_ok=True) # 5) Build and run the inference command upscale_value = upscale_option.replace("x", "") cmd = [ "python", "inference_coz.py", "-i", session_folder, "-o", session_output_dir, "--rec_type", "recursive_multiscale", "--prompt_type", "vlm", "--upscale", upscale_value, "--lora_path", "ckpt/SR_LoRA/model_20001.pkl", "--vae_path", "ckpt/SR_VAE/vae_encoder_20001.pt", "--pretrained_model_name_or_path", "stabilityai/stable-diffusion-3-medium-diffusers", "--ram_ft_path", "ckpt/DAPE/DAPE.pth", "--ram_path", "ckpt/RAM/ram_swin_large_14m.pth" ] try: subprocess.run(cmd, check=True) except subprocess.CalledProcessError as err: print("Inference failed:", err) return [] # 6) Gather output file paths (1.png through 4.png) per_sample_dir = os.path.join(session_output_dir, "per-sample", "input") expected_files = [os.path.join(per_sample_dir, f"{i}.png") for i in range(1, 5)] for fp in expected_files: if not os.path.isfile(fp): print(f"Warning: expected file not found: {fp}") return [] return expected_files def get_caption(src_gallery, evt: gr.SelectData): if not src_gallery or not os.path.isfile(src_gallery[evt.index][0]): return "No caption available." selected_image_path = src_gallery[evt.index][0] base = os.path.basename(selected_image_path) # e.g. "2.png" stem = os.path.splitext(base)[0] # e.g. "2" txt_folder = os.path.join(OUTPUT_DIR, str(evt.index), "per-sample", "input", "txt") txt_path = os.path.join(txt_folder, f"{int(stem) - 1}.txt") if not os.path.isfile(txt_path): return f"Caption file not found: {int(stem) - 1}.txt" try: with open(txt_path, "r", encoding="utf-8") as f: caption = f.read().strip() return caption if caption else "(Caption file is empty.)" except Exception as e: return f"Error reading caption: {e}" css = """ #col-container { margin: 0 auto; max-width: 1024px; } """ with gr.Blocks(css=css) as demo: gr.HTML( """

Chain-of-Zoom

Extreme Super-Resolution via Scale Autoregression and Preference Alignment


""" ) with gr.Column(elem_id="col-container"): with gr.Row(): with gr.Column(): upload_image = gr.Image(label="Upload your input image", type="filepath") upscale_radio = gr.Radio(choices=["1x", "2x", "4x"], value="2x", show_label=False) run_button = gr.Button("Chain-of-Zoom it") preview_with_box = gr.Image(label="Preview (512×512 with centered boxes)", type="pil", interactive=False) with gr.Column(): output_gallery = gr.Gallery(label="Inference Results", show_label=True, columns=[2], rows=[2]) caption_text = gr.Textbox(label="Caption", lines=4, placeholder="Click on any image above to see its caption here.") upload_image.change( fn=lambda img_path, scale_opt: make_preview_with_boxes(img_path, scale_opt) if img_path is not None else None, inputs=[upload_image, upscale_radio], outputs=[preview_with_box] ) upscale_radio.change( fn=lambda img_path, scale_opt: make_preview_with_boxes(img_path, scale_opt) if img_path is not None else None, inputs=[upload_image, upscale_radio], outputs=[preview_with_box] ) # Note: gr.State() will pass session_id automatically run_button.click( fn=run_with_upload, inputs=[upload_image, upscale_radio, gr.State()], outputs=[output_gallery] ) output_gallery.select( fn=get_caption, inputs=[output_gallery], outputs=[caption_text] ) demo.launch(share=True)