File size: 3,980 Bytes
2ae242d |
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 |
import os
import sys
import argparse
import logging
from pathlib import Path
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))
if __name__ == "__main__":
main()
|