wehub-resource-sync
271 lines
9.9 KiB
Python
271 lines
9.9 KiB
Python
#
|
|
# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
#
|
|
|
|
"""Dependency management for TensorRT diffusion demos.
|
|
|
|
Adds the right group's site-packages to sys.path so imports work.
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
|
|
# Marker file to indicate successful installation
|
|
# Must match the marker file used by setup.py
|
|
INSTALL_COMPLETE_MARKER = ".install_complete"
|
|
|
|
# Descriptions for user-facing messages
|
|
GROUP_DESCRIPTIONS = {
|
|
"sd": "SD family (SD 1.4, SDXL, SD3, SD3.5, SVD, Stable Cascade)",
|
|
"flux": "Flux family (Black Forest Labs)",
|
|
"cosmos": "Cosmos family (NVIDIA), Wan2.2 T2V",
|
|
}
|
|
|
|
# Valid dependency groups
|
|
VALID_GROUPS = list(GROUP_DESCRIPTIONS.keys())
|
|
|
|
__all__ = [
|
|
"configure",
|
|
"get_configured_groups",
|
|
"print_status",
|
|
]
|
|
|
|
|
|
def _resolve_deps_root(deps_root: str | None) -> str:
|
|
"""Resolve the dependency root path with precedence: arg > env > default."""
|
|
if deps_root is not None:
|
|
return deps_root
|
|
return os.environ.get("TENSORRT_DIFFUSION_DEPS_ROOT", "/workspace/deps")
|
|
|
|
def _prepend_env_path(var: str, path: str, clean_root: str | None = None):
|
|
"""Prepend `path` to an os.pathsep-separated env var so child processes
|
|
(e.g. the `polygraphy` CLI) inherit the group's dependencies. If
|
|
`clean_root` is given, drop existing entries under it first."""
|
|
def _norm(p: str) -> str:
|
|
return os.path.abspath(os.path.expanduser(p))
|
|
|
|
existing = [p for p in os.environ.get(var, "").split(os.pathsep) if p]
|
|
if clean_root is not None:
|
|
root_abs = _norm(clean_root).rstrip(os.sep) + os.sep
|
|
existing = [p for p in existing if not _norm(p).startswith(root_abs)]
|
|
existing = [p for p in existing if _norm(p) != _norm(path)]
|
|
os.environ[var] = os.pathsep.join([path, *existing])
|
|
|
|
|
|
def _clean_diffusion_paths(deps_root: str = "/workspace/deps"):
|
|
"""Drop any existing paths under deps_root from sys.path."""
|
|
# Filter out any paths that live under deps_root (robust to path forms)
|
|
root_abs = os.path.abspath(os.path.expanduser(deps_root)).rstrip(os.sep) + os.sep
|
|
sys.path[:] = [
|
|
p for p in sys.path
|
|
if not os.path.abspath(os.path.expanduser(p)).startswith(root_abs)
|
|
]
|
|
|
|
|
|
def configure(
|
|
group: str,
|
|
deps_root: str | None = None,
|
|
verbose: bool = False,
|
|
clean: bool = True,
|
|
fallback: bool = False
|
|
):
|
|
"""
|
|
Configure sys.path to use dependencies from the specified group.
|
|
|
|
This function should be called at the top of each demo script, before
|
|
any other imports that depend on external packages.
|
|
|
|
Args:
|
|
group: Dependency group name ("sd", "flux", or "cosmos")
|
|
deps_root: Root directory where dependencies are installed.
|
|
If None, uses TENSORRT_DIFFUSION_DEPS_ROOT environment variable,
|
|
or defaults to "/workspace/deps"
|
|
verbose: Print configuration info (default: False)
|
|
clean: Remove other dependency group paths first (default: True)
|
|
This prevents sys.path inflation in long-running processes.
|
|
fallback: If True, gracefully handle missing dependencies instead of raising
|
|
RuntimeError. Useful for test environments using requirements.txt.
|
|
If False but USE_REQUIREMENTS=1 is set,
|
|
fallback will be automatically enabled.
|
|
|
|
Example:
|
|
from demo_diffusion import deps
|
|
deps.configure("flux")
|
|
|
|
# Or with custom location
|
|
deps.configure("flux", deps_root="/custom/path/deps")
|
|
|
|
# Or via environment variable
|
|
# export TENSORRT_DIFFUSION_DEPS_ROOT=/custom/path/deps
|
|
# deps.configure("flux")
|
|
|
|
# For test environments using requirements.txt (installed via setup.sh)
|
|
# export USE_REQUIREMENTS=1
|
|
# deps.configure("sd") # Will automatically use fallback mode
|
|
"""
|
|
if group not in VALID_GROUPS:
|
|
raise ValueError(
|
|
f"Invalid dependency group: '{group}'\n"
|
|
f"Valid groups: {', '.join(sorted(VALID_GROUPS))}"
|
|
)
|
|
|
|
# Auto-enable fallback mode if using requirements.txt installation
|
|
if not fallback and os.environ.get("USE_REQUIREMENTS") == "1":
|
|
fallback = True
|
|
if verbose:
|
|
print("Detected USE_REQUIREMENTS=1, enabling fallback mode")
|
|
|
|
# Resolve deps_root consistently across the module
|
|
deps_root = _resolve_deps_root(deps_root)
|
|
|
|
# Clean old dependency paths first (prevents inflation)
|
|
if clean:
|
|
_clean_diffusion_paths(deps_root)
|
|
|
|
# Determine Python version
|
|
python_version = f"{sys.version_info.major}.{sys.version_info.minor}"
|
|
|
|
# Construct path to site-packages for this group
|
|
deps_path = os.path.join(
|
|
deps_root,
|
|
group,
|
|
"lib",
|
|
f"python{python_version}",
|
|
"site-packages"
|
|
)
|
|
|
|
# Check if installation is complete
|
|
group_dir = os.path.join(deps_root, group)
|
|
marker_file = os.path.join(group_dir, INSTALL_COMPLETE_MARKER)
|
|
|
|
if os.path.exists(deps_path) and os.path.exists(marker_file):
|
|
# Insert at the beginning to override any system packages
|
|
sys.path.insert(0, deps_path)
|
|
|
|
# Mirror onto the environment so child processes (e.g. the polygraphy
|
|
# CLI that engine builds shell out to) use these deps too. When clean,
|
|
# drop other groups' entries under deps_root, matching sys.path above.
|
|
clean_root = deps_root if clean else None
|
|
_prepend_env_path("PYTHONPATH", deps_path, clean_root=clean_root)
|
|
_prepend_env_path("PATH", os.path.join(group_dir, "bin"), clean_root=clean_root)
|
|
|
|
if verbose:
|
|
description = GROUP_DESCRIPTIONS.get(group, group)
|
|
print(f"Configured dependencies: {description}")
|
|
print(f" Path: {deps_path}")
|
|
else:
|
|
# Dependencies not found or installation incomplete
|
|
description = GROUP_DESCRIPTIONS.get(group, group)
|
|
|
|
# Check if it's an incomplete installation
|
|
if os.path.exists(group_dir) and not os.path.exists(marker_file):
|
|
error_msg = (
|
|
f"Dependencies for '{group}' are incomplete!\n"
|
|
f" Location: {group_dir}\n"
|
|
f" Description: {description}\n"
|
|
f" A previous installation failed or was interrupted.\n"
|
|
f" To fix, run:\n"
|
|
f" python setup.py {group}\n"
|
|
f" (This will clean up and reinstall)"
|
|
)
|
|
if fallback:
|
|
if verbose:
|
|
print(f"Warning: {error_msg}")
|
|
print("Continuing with fallback mode...")
|
|
return
|
|
else:
|
|
raise RuntimeError(error_msg)
|
|
else:
|
|
# Not installed at all
|
|
error_msg = (
|
|
f"Dependencies for '{group}' not found!\n"
|
|
f" Expected at: {deps_path}\n"
|
|
f" Description: {description}\n"
|
|
f" To install, run: python setup.py {group}\n"
|
|
f" If you installed elsewhere, set TENSORRT_DIFFUSION_DEPS_ROOT or pass deps_root."
|
|
)
|
|
if fallback:
|
|
if verbose:
|
|
print(f"Warning: {error_msg}")
|
|
print("Continuing with fallback mode (assuming requirements.txt setup)...")
|
|
return
|
|
else:
|
|
raise RuntimeError(error_msg)
|
|
|
|
|
|
def get_configured_groups(deps_root: str | None = None) -> list[str]:
|
|
"""
|
|
Get list of dependency groups that are currently installed.
|
|
|
|
A group is considered installed only if both:
|
|
1. The group directory exists
|
|
2. The .install_complete marker file exists in that directory
|
|
|
|
Args:
|
|
deps_root: Root directory where dependencies are installed.
|
|
If None, uses TENSORRT_DIFFUSION_DEPS_ROOT environment variable,
|
|
or defaults to "/workspace/deps"
|
|
|
|
Returns:
|
|
List of installed group names (e.g., ["sd", "flux"])
|
|
"""
|
|
# Resolve deps_root
|
|
deps_root = _resolve_deps_root(deps_root)
|
|
|
|
installed = []
|
|
for group in VALID_GROUPS:
|
|
# Check if the group directory exists and has the completion marker
|
|
group_dir = os.path.join(deps_root, group)
|
|
marker_file = os.path.join(group_dir, INSTALL_COMPLETE_MARKER)
|
|
|
|
# Only consider it installed if marker file exists
|
|
if os.path.exists(group_dir) and os.path.isdir(group_dir) and os.path.exists(marker_file):
|
|
installed.append(group)
|
|
return sorted(installed)
|
|
|
|
|
|
def print_status(deps_root: str | None = None):
|
|
"""
|
|
Print the status of all dependency groups.
|
|
|
|
Args:
|
|
deps_root: Root directory where dependencies are installed.
|
|
If None, uses TENSORRT_DIFFUSION_DEPS_ROOT environment variable,
|
|
or defaults to "/workspace/deps"
|
|
"""
|
|
# Resolve deps_root
|
|
deps_root = _resolve_deps_root(deps_root)
|
|
|
|
print("Dependency Groups Status:")
|
|
print("-" * 60)
|
|
print(f"Location: {deps_root}")
|
|
print("-" * 60)
|
|
|
|
installed = get_configured_groups(deps_root)
|
|
|
|
for group in sorted(VALID_GROUPS):
|
|
description = GROUP_DESCRIPTIONS.get(group, group)
|
|
status = "Installed" if group in installed else "Not installed"
|
|
print(f" {group:8} - {status:15} - {description}")
|
|
|
|
print("-" * 60)
|
|
|
|
if not installed:
|
|
print("No dependency groups installed.")
|
|
print("Run: python setup.py all")
|
|
else:
|
|
print(f"{len(installed)} group(s) installed: {', '.join(installed)}")
|