"""Asset retrieval route functions for ComfyUI API."""
import httpx
from src.auth.base import ComfyAuth
from src.client.get_data import get_data
from src.client.response import ResponseGetData
from src.utils.logging import log_call
class AssetError(Exception):
"""Raised when asset operations fail."""
def __init__(self, message: str, response: ResponseGetData | None = None):
super().__init__(message)
self.response = response
@log_call(action_name="get_asset", level_name="route")
async def get_asset(
auth: ComfyAuth,
filename: str,
subfolder: str = "",
asset_type: str = "output",
*,
session: httpx.AsyncClient | None = None,
) -> ResponseGetData:
"""Retrieve a generated asset (image, audio, etc.) from ComfyUI.
Downloads the binary content of a generated asset using the /view endpoint.
The asset is returned in the response body as bytes.
Args:
auth: ComfyAuth instance for base URL and authentication
filename: Asset filename (e.g., "ComfyUI_00001.png")
subfolder: Optional subfolder within output directory
asset_type: Asset type ("output", "input", or "temp")
session: Optional HTTPX client for connection pooling
Returns:
ResponseGetData with:
- response: Binary asset content (bytes)
- status: 200 on success
- is_success: True on success
Raises:
AssetError: If asset retrieval fails
Example:
>>> auth = NoAuth("http://127.0.0.1:8188")
>>> res = await get_asset(
... auth=auth,
... filename="ComfyUI_00001.png",
... subfolder="",
... asset_type="output",
... )
>>> print(type(res.response), len(res.response))
<class 'bytes'> 245678
"""
url = f"{auth.base_url}/view"
params = {
"filename": filename,
"type": asset_type,
}
if subfolder:
params["subfolder"] = subfolder
res = await get_data(
auth=auth,
method="GET",
url=url,
params=params,
session=session,
)
if not res.is_success:
raise AssetError(
f"Failed to get asset {filename}: HTTP {res.status}",
response=res,
)
return res
@log_call(action_name="get_asset_metadata", level_name="route", log_level="DEBUG")
async def get_asset_metadata(
auth: ComfyAuth,
filename: str,
subfolder: str = "",
asset_type: str = "output",
*,
session: httpx.AsyncClient | None = None,
) -> ResponseGetData:
"""Retrieve metadata for an asset using HEAD request.
Returns HTTP headers for an asset without downloading the full content.
Useful for checking Content-Type, Content-Length, etc.
Args:
auth: ComfyAuth instance for base URL and authentication
filename: Asset filename (e.g., "ComfyUI_00001.png")
subfolder: Optional subfolder within output directory
asset_type: Asset type ("output", "input", or "temp")
session: Optional HTTPX client for connection pooling
Returns:
ResponseGetData with:
- response: None (HEAD request has no body)
- headers: Asset HTTP headers
- status: 200 on success
- is_success: True on success
Raises:
AssetError: If metadata retrieval fails
Example:
>>> auth = NoAuth("http://127.0.0.1:8188")
>>> res = await get_asset_metadata(
... auth=auth,
... filename="ComfyUI_00001.png",
... )
>>> print(res.headers["content-type"])
"image/png"
"""
url = f"{auth.base_url}/view"
params = {
"filename": filename,
"type": asset_type,
}
if subfolder:
params["subfolder"] = subfolder
res = await get_data(
auth=auth,
method="HEAD",
url=url,
params=params,
session=session,
)
if not res.is_success:
raise AssetError(
f"Failed to get asset metadata for {filename}: HTTP {res.status}",
response=res,
)
return res
