Spaces:
Running
on
Zero
Running
on
Zero
File size: 5,196 Bytes
0301e15 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 |
import gradio as gr
import subprocess
import os
import shutil
from pathlib import Path
from PIL import Image
import spaces
# -----------------------------------------------------------------------------
# CONFIGURE THESE PATHS TO MATCH YOUR PROJECT STRUCTURE
# -----------------------------------------------------------------------------
INPUT_DIR = "samples"
OUTPUT_DIR = "inference_results/coz_vlmprompt"
# -----------------------------------------------------------------------------
# HELPER FUNCTION TO RUN INFERENCE AND RETURN THE OUTPUT IMAGE
# -----------------------------------------------------------------------------
@spaces.GPU()
def run_with_upload(uploaded_image_path):
"""
1) Clear out INPUT_DIR (so old samples don’t linger).
2) Copy the uploaded image into INPUT_DIR.
3) Run your inference_coz.py command (which reads from -i INPUT_DIR).
4) After it finishes, find the most recently‐modified PNG in OUTPUT_DIR.
5) Return a PIL.Image, which Gradio will display.
"""
# 1) Make sure INPUT_DIR exists; if it does, delete everything inside.
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}")
# 2) Copy the uploaded image into INPUT_DIR.
# Gradio will give us a path like "/tmp/gradio_xyz.png"
if uploaded_image_path is None:
return None
try:
# Open with PIL (this handles JPEG, BMP, TIFF, etc.)
pil_img = Image.open(uploaded_image_path).convert("RGB")
except Exception as e:
print(f"Error: could not open uploaded image: {e}")
return None
# Save it as "input.png" in our INPUT_DIR
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 None
# 3) Build and run your inference_coz.py command.
# This will block until it completes.
cmd = [
"python", "inference_coz.py",
"-i", INPUT_DIR,
"-o", OUTPUT_DIR,
"--rec_type", "recursive_multiscale",
"--prompt_type", "vlm",
"--upscale", "2",
"--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:
# If inference_coz.py crashes, we can print/log the error.
print("Inference failed:", err)
return None
# 4) After it finishes, scan OUTPUT_DIR for .png files.
RECUSIVE_DIR = f'{OUTPUT_DIR}/recursive'
if not os.path.isdir(RECUSIVE_DIR):
return None
png_files = [
os.path.join(RECUSIVE_DIR, fn)
for fn in os.listdir(RECUSIVE_DIR)
if fn.lower().endswith(".png")
]
if not png_files:
return None
# 5) Pick the most recently‐modified PNG
latest_png = max(png_files, key=os.path.getmtime)
# 6) Open and return a PIL.Image. Gradio will display it automatically.
try:
img = Image.open(latest_png).convert("RGB")
except Exception as e:
print(f"Error opening {latest_png}: {e}")
return None
return img
# -----------------------------------------------------------------------------
# BUILD THE GRADIO INTERFACE
# -----------------------------------------------------------------------------
with gr.Blocks() as demo:
gr.Markdown("## Upload an image, then click **Run Inference** to process it.")
# 1) Image upload component. We set type="filepath" so the callback
# (run_with_upload) will receive a local path to the uploaded file.
upload_image = gr.Image(
label="Upload your input image",
type="filepath"
)
# 2) A button that the user will click to launch inference.
run_button = gr.Button("Run Inference")
# 3) An output <Image> where we will show the final PNG.
output_image = gr.Image(
label="Inference Result",
type="pil" # because run_with_upload() returns a PIL.Image
)
# Wire the button: when clicked, call run_with_upload(upload_image), put
# its return value into output_image.
run_button.click(
fn=run_with_upload,
inputs=upload_image,
outputs=output_image
)
# -----------------------------------------------------------------------------
# START THE GRADIO SERVER
# -----------------------------------------------------------------------------
demo.launch(share=True)
|