Move InvSR download to runtime (avoids build OOM), add status endpoint

Dockerfile

@@ -9,25 +9,11 @@ USER user
 ENV PATH="/home/user/.local/bin:$PATH"
 WORKDIR /app
 
-# Clone InvSR source
 RUN git clone --depth 1 https://github.com/zsyOAOA/InvSR.git /app/InvSR
 
-# Install pip deps
 COPY --chown=user requirements.txt .
 RUN pip install --no-cache-dir -r requirements.txt
 
-# Pre-download SD-Turbo (~5GB) + noise predictor at BUILD time
-RUN python -c "\
-from huggingface_hub import snapshot_download;\
-snapshot_download('stabilityai/sd-turbo', cache_dir='/app/models/invsr');\
-print('SD-Turbo downloaded')\
-"
-RUN python -c "\
-from huggingface_hub import hf_hub_download;\
-hf_hub_download('OAOA/InvSR', 'noise_predictor_sd_turbo_v5.pth', cache_dir='/app/models/invsr');\
-print('Noise predictor downloaded')\
-"
-
 COPY --chown=user app.py .
 
 CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860", "--timeout-keep-alive", "600"]

app.py

@@ -78,81 +78,58 @@ def upscale_mewzoom(image_bytes: bytes, scale: str) -> tuple[bytes, dict]:
 # ── InvSR Model (Diffusion 4X) ──────────────────────────────
 _INVSR_PATH = Path("/app/InvSR")
 _sampler_invsr = None
+_invsr_status = "not_loaded"
+_invsr_error = None
 
 
-def _patch_invsr():
-    """Patch InvSR source for CPU/float32 support."""
-    p = _INVSR_PATH / "sampler_invsr.py"
-    code = p.read_text()
-
-    # Remove basicsr import chain (not needed for inference)
-    code = code.replace("from datapipe.datasets import create_dataset", "")
-
-    # Add device param to BaseSampler
-    old_init = """class BaseSampler:
-    def __init__(self, configs):
-        '''
-        Input:
-            configs: config, see the yaml file in folder ./configs/
-                configs.sampler_config.{start_timesteps, padding_mod, seed, sf, num_sample_steps}
-            seed: int, random seed
-        '''
-        self.configs = configs
-
-        self.setup_seed()
-
-        self.build_model()
-
-    def setup_seed(self, seed=None):
-        seed = self.configs.seed if seed is None else seed
-        random.seed(seed)
-        np.random.seed(seed)
-        torch.manual_seed(seed)
-        torch.cuda.manual_seed_all(seed)"""
-
-    new_init = """class BaseSampler:
-    def __init__(self, configs, device='auto'):
-        self.configs = configs
-        if device == 'auto':
-            device = 'cuda' if torch.cuda.is_available() else 'cpu'
-        self.device = torch.device(device)
-        self.dtype = torch.float16 if self.device.type == 'cuda' else torch.float32
-        self.setup_seed()
-        self.build_model()
-
-    def setup_seed(self, seed=None):
-        seed = self.configs.seed if seed is None else seed
-        random.seed(seed)
-        np.random.seed(seed)
-        torch.manual_seed(seed)
-        if torch.cuda.is_available():
-            torch.cuda.manual_seed_all(seed)"""
-
-    code = code.replace(old_init, new_init)
-
-    # Replace .cuda() and .type(torch.float16) with device-aware versions
-    code = code.replace('sd_pipe.to(f"cuda")', "sd_pipe.to(self.device)")
-    code = code.replace("model_start.cuda()", "model_start.to(self.device)")
-    code = code.replace('map_location=f"cuda"', "map_location=self.device")
-    code = code.replace("im_cond.type(torch.float16)", "im_cond.type(self.dtype)")
-    code = code.replace(".type(torch.float16)", ".type(self.dtype)")
-    code = code.replace("data['lq'].cuda()", "data['lq'].to(self.device)")
-    code = code.replace("util_image.img2tensor(im_cond).cuda()", "util_image.img2tensor(im_cond).to(self.device)")
-
-    # Lazy import create_dataset in inference method
-    code = code.replace(
-        "if in_path.is_dir():\n            data_config",
-        "if in_path.is_dir():\n            from datapipe.datasets import create_dataset\n            data_config",
-    )
-
-    p.write_text(code)
-    logger.info("InvSR sampler patched for CPU/float32")
-
-
-def _load_invsr():
-    global _sampler_invsr
-    if _sampler_invsr is not None:
-        return _sampler_invsr
+def _load_invsr_sync():
+    """Download + load InvSR (called in background during startup)"""
+    global _sampler_invsr, _invsr_status, _invsr_error
+    try:
+        _invsr_status = "downloading"
+        _patch_invsr()
+        sys.path.insert(0, str(_INVSR_PATH))
+        sys.path.insert(0, str(_INVSR_PATH / "src"))
+
+        from omegaconf import OmegaConf
+        from huggingface_hub import snapshot_download, hf_hub_download
+
+        invsr_cache = str(CACHE_DIR / "invsr")
+        CACHE_DIR.mkdir(exist_ok=True)
+
+        logger.info("Downloading SD-Turbo (~5GB, one-time)...")
+        snapshot_download("stabilityai/sd-turbo", cache_dir=invsr_cache, resume_download=True)
+        logger.info("SD-Turbo downloaded")
+
+        logger.info("Downloading noise predictor...")
+        hf_hub_download("OAOA/InvSR", "noise_predictor_sd_turbo_v5.pth", cache_dir=invsr_cache)
+        ckpt = None
+        for f in Path(invsr_cache).rglob("noise_predictor_sd_turbo_v5.pth"):
+            ckpt = str(f); break
+        if not ckpt:
+            raise FileNotFoundError("Noise predictor not found after download")
+
+        _invsr_status = "loading"
+        from sampler_invsr import InvSamplerSR
+        cfg = OmegaConf.load(str(_INVSR_PATH / "configs" / "sample-sd-turbo.yaml"))
+        cfg.sd_pipe.params.torch_dtype = "torch.float32" if _DEVICE == "cpu" else "torch.float16"
+        cfg.sd_pipe.params.cache_dir = invsr_cache
+        cfg.sd_pipe.params.local_files_only = True
+        cfg.model_start.ckpt_path = ckpt
+        cfg.timesteps = [200]; cfg.bs = 1; cfg.tiled_vae = True
+        cfg.color_fix = "wavelet"; cfg.basesr.chopping.pch_size = 128
+        cfg.basesr.chopping.extra_bs = 8
+
+        logger.info("Loading InvSR into memory...")
+        _sampler_invsr = InvSamplerSR(cfg, device="auto")
+        if _DEVICE == "cpu":
+            _sampler_invsr.sd_pipe = _sampler_invsr.sd_pipe.to(dtype=torch.float32)
+        _invsr_status = "ready"
+        logger.info("InvSR ready on %s", _DEVICE)
+    except Exception as e:
+        _invsr_status = "error"
+        _invsr_error = str(e)
+        logger.error("InvSR load failed: %s", e)
 
     _patch_invsr()
     sys.path.insert(0, str(_INVSR_PATH))
@@ -190,7 +167,12 @@ def _load_invsr():
 
 
 def upscale_invsr(image_bytes: bytes, num_steps: int = 1) -> bytes:
-    sampler = _load_invsr()
+    global _sampler_invsr
+    if _invsr_status == "error":
+        raise HTTPException(500, f"InvSR failed to load: {_invsr_error}")
+    if _sampler_invsr is None:
+        raise HTTPException(503, f"InvSR is {_invsr_status}. Check /health for status.")
+    sampler = _sampler_invsr
     sys.path.insert(0, str(_INVSR_PATH))
     from utils import util_image
 
@@ -271,6 +253,9 @@ async def lifespan(app: FastAPI):
     logger.info("Loading MewZoom models...")
     for s in MEWZOOM_MODELS:
         _load_mewzoom(s)
+    # Start InvSR download+load in background thread
+    import threading
+    threading.Thread(target=_load_invsr_sync, daemon=True).start()
     yield
 
 
@@ -286,7 +271,14 @@ app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], all
 @app.get("/")
 @app.get("/health")
 async def health():
-    return JSONResponse({"status": "healthy", "device": _DEVICE, "models": ["2x", "4x", "invsr"], "gpu": torch.cuda.is_available()})
+    return JSONResponse({
+        "status": "healthy",
+        "device": _DEVICE,
+        "models": list(MEWZOOM_MODELS.keys()) + ["invsr"],
+        "gpu": torch.cuda.is_available(),
+        "invsr_status": _invsr_status,
+        "invsr_error": _invsr_error,
+    })
 
 
 @app.post("/upscale/2x")