Update app.py

app.py

@@ -3,9 +3,11 @@ import subprocess
 import os
 import shutil
 from pathlib import Path
+from inference_coz_single import recursive_multiscale_sr
 from PIL import Image, ImageDraw
 import spaces
 
+
 # ------------------------------------------------------------------
 # CONFIGURE THESE PATHS TO MATCH YOUR PROJECT STRUCTURE
 # ------------------------------------------------------------------
@@ -85,102 +87,13 @@ def make_preview_with_boxes(image_path: str, scale_option: str) -> Image.Image:
     return base
 
 
-# ------------------------------------------------------------------
-# HELPER FUNCTIONS FOR INFERENCE & CAPTION (unchanged from your original)
-# ------------------------------------------------------------------
 @spaces.GPU(duration=120)
 def run_with_upload(uploaded_image_path, upscale_option):
-    """
-    1) Clear INPUT_DIR  
-    2) Save the uploaded file as input.png in INPUT_DIR  
-    3) Read `upscale_option` (e.g. "1x", "2x", "4x") → turn it into "1","2","4"  
-    4) Call inference_coz.py with `--upscale <that_value>`  
-    5) Return the FOUR output‐PNG file‐paths as a Python list, so that Gradio's Gallery
-       can display them.
-    """
-    # ————————————————————————————————————————————————————————————
-    # (Copy‐paste exactly your existing code here; no changes needed)
-    # ————————————————————————————————————————————————————————————
-
-    os.makedirs(INPUT_DIR, exist_ok=True)
-    for fn in os.listdir(INPUT_DIR):
-        full_path = os.path.join(INPUT_DIR, fn)
-        try:
-            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)
-        except Exception as e:
-            print(f"Warning: could not delete {full_path}: {e}")
-
-    if uploaded_image_path is None:
-        return []
-    try:
-        pil_img = Image.open(uploaded_image_path).convert("RGB")
-    except Exception as e:
-        print(f"Error: could not open uploaded image: {e}")
-        return []
-    save_path = Path(INPUT_DIR) / "input.png"
-    try:
-        pil_img.save(save_path, format="PNG")
-    except Exception as e:
-        print(f"Error: could not save as PNG: {e}")
-        return []
-
-    upscale_value = upscale_option.replace("x", "")  # e.g. "2x" → "2"
-    cmd = [
-        "python", "inference_coz.py",
-        "-i", INPUT_DIR,
-        "-o", 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 []
-
-    per_sample_dir = os.path.join(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):
-    """
-    Given a clicked‐on image in the gallery, read the corresponding .txt in
-    .../per-sample/input/txt and return its contents.
-    """
-    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, "per-sample", "input", "txt")
-    txt_path = os.path.join(txt_folder, f"{int(stem) - 1}.txt")
+  upscale_value = upscale_option.replace("x", "")  # e.g. "2x" → "2"
+  
+  return recursive_multiscale_sr(uploaded_image_path, int(upscale_value))[0]
 
-    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}"
 
 
 # ------------------------------------------------------------------
@@ -248,13 +161,6 @@ with gr.Blocks(css=css) as demo:
                     columns=[2], rows=[2]
                 )
 
-                # 6) Textbox under the gallery for showing captions
-                caption_text = gr.Textbox(
-                    label="Caption",
-                    lines=4,
-                    placeholder="Click on any image above to see its caption here."
-                )
-
         # ------------------------------------------------------------------
         # CALLBACK #1: Whenever the user uploads or changes the radio, update preview
         # ------------------------------------------------------------------
@@ -292,23 +198,10 @@ with gr.Blocks(css=css) as demo:
             outputs=[output_gallery]
         )
 
-        # ------------------------------------------------------------------
-        # CALLBACK #3: When an image in the gallery is clicked, show its caption
-        # ------------------------------------------------------------------
-
-        output_gallery.select(
-            fn=get_caption,
-            inputs=[output_gallery],
-            outputs=[caption_text]
-        )
 
 # ------------------------------------------------------------------
 # START THE GRADIO SERVER
 # ------------------------------------------------------------------
 
-# 🔧 1) turn the global queue ON and set its default_concurrency_limit to 1  
-demo.queue(default_concurrency_limit=1,        # ≤ 1 worker per event
-           max_size=20)                        # optional: allow 20 waiting jobs
-
 # 🔧 2) launch as usual
 demo.launch(share=True)