update Postman collection

app/colorize_model.py

@@ -33,11 +33,22 @@ class ColorizeModel:
         self.hf_token = os.getenv("HF_TOKEN") or None
 
         # Configure writable cache to avoid permission issues on Spaces
-        hf_cache_dir = os.getenv("HF_HOME", "./hf_cache")
+        # Prefer user home cache: ~/.cache/huggingface
+        default_home_cache = os.path.join(os.path.expanduser("~"), ".cache", "huggingface")
+        hf_cache_dir = os.getenv("HF_HOME", default_home_cache)
         os.environ.setdefault("HF_HOME", hf_cache_dir)
         os.environ.setdefault("HUGGINGFACE_HUB_CACHE", hf_cache_dir)
         os.environ.setdefault("TRANSFORMERS_CACHE", hf_cache_dir)
-        os.makedirs(hf_cache_dir, exist_ok=True)
+        try:
+            os.makedirs(hf_cache_dir, exist_ok=True)
+        except Exception:
+            # Fallback to a local data dir if home is not writable
+            hf_cache_dir = os.path.abspath(os.path.join(".", "data", "hf_cache"))
+            os.environ["HF_HOME"] = hf_cache_dir
+            os.environ["HUGGINGFACE_HUB_CACHE"] = hf_cache_dir
+            os.environ["TRANSFORMERS_CACHE"] = hf_cache_dir
+            os.makedirs(hf_cache_dir, exist_ok=True)
+        logger.info("HF cache directory: %s", hf_cache_dir)
 
         # Avoid libgomp warning by setting a valid integer
         os.environ.setdefault("OMP_NUM_THREADS", "1")

app/main.py

@@ -61,15 +61,23 @@ else:
     except:
         pass
 
-# Create directories
-UPLOAD_DIR = Path("uploads")
-RESULT_DIR = Path("results")
-UPLOAD_DIR.mkdir(exist_ok=True)
-RESULT_DIR.mkdir(exist_ok=True)
+# Create writable directories (prefer under user home unless overridden)
+def _resolve_dir(path_str: str, default_subdir: str) -> Path:
+    p = Path(path_str).expanduser()
+    if not p.is_absolute():
+        # Place relative paths under ~/data to avoid permission issues
+        p = Path.home() / "data" / default_subdir
+    return p
 
-# Mount static files for serving results
-app.mount("/results", StaticFiles(directory="results"), name="results")
-app.mount("/uploads", StaticFiles(directory="uploads"), name="uploads")
+UPLOAD_DIR = _resolve_dir(os.getenv("UPLOAD_DIR", settings.UPLOAD_DIR), "uploads")
+RESULT_DIR = _resolve_dir(os.getenv("RESULT_DIR", settings.RESULT_DIR), "results")
+UPLOAD_DIR.mkdir(parents=True, exist_ok=True)
+RESULT_DIR.mkdir(parents=True, exist_ok=True)
+logger.info("Storage directories -> uploads: %s, results: %s", str(UPLOAD_DIR), str(RESULT_DIR))
+
+# Mount static files for serving results from resolved directories
+app.mount("/results", StaticFiles(directory=str(RESULT_DIR)), name="results")
+app.mount("/uploads", StaticFiles(directory=str(UPLOAD_DIR)), name="uploads")
 
 # Initialize ColorizeNet model
 colorize_model = None

test.py

@@ -1,134 +1,162 @@
-import os
-import sys
 import argparse
+import io
+import json
 import logging
-from pathlib import Path
+import os
+import sys
+import time
+from typing import Optional
 
 import requests
 
 
-logging.basicConfig(
-    level=logging.INFO,
-    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
-)
-logger = logging.getLogger("colorize_test")
-
-
-def normalize_space_url(url: str) -> str:
-    """
-    Accept either a runtime URL (https://<owner>-<space>.hf.space)
-    or a repo URL (https://huggingface.co/spaces/<Owner>/<Space>?logs=container)
-    and return a runtime base URL suitable for API calls.
-    """
-    if not url:
-        return url
-    url = url.strip()
-    # Strip query params
-    if "?" in url:
-        url = url.split("?", 1)[0]
-    # Already runtime URL
-    if ".hf.space" in url:
-        return url.rstrip("/")
-    # Convert repo URL to runtime URL
-    marker = "/spaces/"
-    if "huggingface.co" in url and marker in url:
-        try:
-            _, rest = url.split(marker, 1)
-            owner, space = rest.split("/", 1)
-            runtime = f"https://{owner.lower()}-{space.lower()}.hf.space"
-            return runtime.rstrip("/")
-        except Exception:
-            return url.rstrip("/")
-    return url.rstrip("/")
-
-
-def main() -> None:
-    parser = argparse.ArgumentParser(description="Test Colorize API")
-    parser.add_argument(
-        "--image",
-        "-i",
-        type=str,
-        required=False,
-        default=r"C:\projects\colorize_text\pexels-andrey-grushnikov-223358-707676.jpg",
-        help="Path to input image file",
-    )
-    parser.add_argument(
-        "--timeout",
-        type=int,
-        default=180,
-        help="Request timeout in seconds (default: 180)",
-    )
-    args = parser.parse_args()
-
-    base_url = os.getenv(
-        "BASE_URL",
-        os.getenv(
-            "SPACE_HOST",
-            "http://localhost:7860",
-        ),
-    )
-    base_url = normalize_space_url(base_url)
-    app_check_token = os.getenv("APP_CHECK_TOKEN", "").strip()
-
-    image_path = Path(args.image)
-    if not image_path.exists():
-        logger.error("Image file not found at: %s", str(image_path))
-        sys.exit(1)
-
-    # Health check
-    try:
-        logger.info("Using BASE_URL: %s", base_url)
-        health = requests.get(f"{base_url}/health", timeout=15)
-        logger.info("Health check %s: %s", health.status_code, health.text)
-    except Exception as e:
-        logger.warning("Health check failed: %s", str(e))
-
-    headers = {}
-    if app_check_token:
-        headers["X-Firebase-AppCheck"] = app_check_token
-
-    files = {"file": (image_path.name, open(image_path, "rb"), "application/octet-stream")}
-
-    logger.info("Sending request to %s", f"{base_url}/colorize")
-    try:
-        resp = requests.post(
-            f"{base_url}/colorize",
-            headers=headers,
-            files=files,
-            timeout=args.timeout,
-        )
-    finally:
-        files["file"][1].close()
-
-    logger.info("Response status: %s", resp.status_code)
-    if not resp.ok:
-        logger.error("Error response: %s", resp.text)
-        sys.exit(2)
-
-    data = resp.json()
-    download_url = data.get("download_url")
-    api_download_url = data.get("api_download_url")
-    filename = data.get("filename")
-
-    logger.info("Colorization success: %s", data.get("success"))
-    logger.info("Download URL: %s", download_url)
-    logger.info("API Download URL: %s", api_download_url)
-    logger.info("Filename: %s", filename)
-
-    # Optionally download the result locally
-    if download_url:
-        try:
-            out_path = Path("colorized_result.jpg")
-            r = requests.get(download_url, timeout=60)
-            if r.ok:
-                out_path.write_bytes(r.content)
-                logger.info("Saved colorized image to: %s", str(out_path.resolve()))
-            else:
-                logger.warning("Could not download from public URL, status: %s", r.status_code)
-        except Exception as e:
-            logger.warning("Download via public URL failed: %s", str(e))
+def configure_logging(verbose: bool) -> None:
+	log_level = logging.DEBUG if verbose else logging.INFO
+	logging.basicConfig(
+		level=log_level,
+		format="%(asctime)s - %(levelname)s - %(message)s"
+	)
+
+
+def get_base_url(cli_base_url: Optional[str]) -> str:
+	if cli_base_url:
+		return cli_base_url.rstrip("/")
+	env_base = os.getenv("BASE_URL")
+	if env_base:
+		return env_base.rstrip("/")
+	# Fallback to HF style URL if provided via POSTMAN collection, else localhost
+	return "http://localhost:7860"
+
+
+def wait_for_model(base_url: str, timeout_seconds: int = 300) -> None:
+	deadline = time.time() + timeout_seconds
+	health_url = f"{base_url}/health"
+	logging.info("Waiting for model to load at %s", health_url)
+	last_status = None
+	while time.time() < deadline:
+		try:
+			resp = requests.get(health_url, timeout=15)
+			if resp.ok:
+				data = resp.json()
+				last_status = data
+				if data.get("model_loaded"):
+					logging.info("Model loaded: %s", json.dumps(data))
+					return
+				logging.info("Health: %s", json.dumps(data))
+			else:
+				logging.warning("Health check HTTP %s", resp.status_code)
+		except Exception as e:
+			logging.warning("Health check error: %s", str(e))
+		time.sleep(3)
+	raise RuntimeError("Model did not load before timeout. Last health: %s" % (last_status,))
+
+
+def upload_image(base_url: str, image_path: str, auth_bearer: Optional[str], app_check: Optional[str]) -> dict:
+	url = f"{base_url}/upload"
+	headers = {}
+	if auth_bearer:
+		headers["Authorization"] = f"Bearer {auth_bearer}"
+	if app_check:
+		headers["X-Firebase-AppCheck"] = app_check
+	with open(image_path, "rb") as f:
+		files = {"file": (os.path.basename(image_path), f, "image/jpeg")}
+		resp = requests.post(url, files=files, headers=headers, timeout=120)
+	if not resp.ok:
+		raise RuntimeError("Upload failed: HTTP %s %s" % (resp.status_code, resp.text))
+	data = resp.json()
+	logging.info("Upload response: %s", json.dumps(data))
+	return data
+
+
+def colorize_image(base_url: str, image_path: str, auth_bearer: Optional[str], app_check: Optional[str]) -> dict:
+	url = f"{base_url}/colorize"
+	headers = {}
+	if auth_bearer:
+		headers["Authorization"] = f"Bearer {auth_bearer}"
+	if app_check:
+		headers["X-Firebase-AppCheck"] = app_check
+	with open(image_path, "rb") as f:
+		files = {"file": (os.path.basename(image_path), f, "image/jpeg")}
+		resp = requests.post(url, files=files, headers=headers, timeout=900)
+	if not resp.ok:
+		raise RuntimeError("Colorize failed: HTTP %s %s" % (resp.status_code, resp.text))
+	data = resp.json()
+	logging.info("Colorize response: %s", json.dumps(data))
+	return data
+
+
+def download_result(base_url: str, result_id: str, output_path: str, auth_bearer: Optional[str], app_check: Optional[str]) -> None:
+	url = f"{base_url}/download/{result_id}"
+	headers = {}
+	if auth_bearer:
+		headers["Authorization"] = f"Bearer {auth_bearer}"
+	if app_check:
+		headers["X-Firebase-AppCheck"] = app_check
+	resp = requests.get(url, headers=headers, stream=True, timeout=300)
+	if not resp.ok:
+		raise RuntimeError("Download failed: HTTP %s %s" % (resp.status_code, resp.text))
+	with open(output_path, "wb") as out:
+		for chunk in resp.iter_content(chunk_size=8192):
+			if chunk:
+				out.write(chunk)
+	logging.info("Saved colorized image to: %s", output_path)
+
+
+def main() -> int:
+	parser = argparse.ArgumentParser(description="End-to-end test for Colorize API")
+	parser.add_argument("--base-url", type=str, help="API base URL, e.g. https://<space>.hf.space")
+	parser.add_argument("--image", type=str, required=True, help="Path to input image")
+	parser.add_argument("--out", type=str, default="colorized_result.jpg", help="Path to save colorized image")
+	parser.add_argument("--auth", type=str, default=os.getenv("ID_TOKEN", ""), help="Optional Firebase id_token")
+	parser.add_argument("--app-check", type=str, default=os.getenv("APP_CHECK_TOKEN", ""), help="Optional App Check token")
+	parser.add_argument("--skip-wait", action="store_true", help="Skip waiting for model to load")
+	parser.add_argument("--verbose", action="store_true", help="Verbose logging")
+	args = parser.parse_args()
+
+	configure_logging(args.verbose)
+	base_url = get_base_url(args.base_url)
+	image_path = args.image
+
+	if not os.path.exists(image_path):
+		logging.error("Image not found: %s", image_path)
+		return 1
+
+	if not args.skip_wait:
+		try:
+			wait_for_model(base_url, timeout_seconds=600)
+		except Exception as e:
+			logging.warning("Continuing despite health wait failure: %s", str(e))
+
+	auth_bearer = args.auth.strip() or None
+	app_check = args.app_check.strip() or None
+
+	try:
+		upload_resp = upload_image(base_url, image_path, auth_bearer, app_check)
+	except Exception as e:
+		logging.error("Upload error: %s", str(e))
+		return 1
+
+	try:
+		colorize_resp = colorize_image(base_url, image_path, auth_bearer, app_check)
+	except Exception as e:
+		logging.error("Colorize error: %s", str(e))
+		return 1
+
+	result_id = colorize_resp.get("result_id")
+	if not result_id:
+		logging.error("No result_id in response: %s", json.dumps(colorize_resp))
+		return 1
+
+	try:
+		download_result(base_url, result_id, args.out, auth_bearer, app_check)
+	except Exception as e:
+		logging.error("Download error: %s", str(e))
+		return 1
+
+	logging.info("Test workflow completed successfully.")
+	return 0
 
 
 if __name__ == "__main__":
-    main()
+	sys.exit(main())