[diffusion] multi-platform: support diffusion on amd and fix encoder loading on MI325 (#13760)

Co-authored-by: Sabre Shao <sabre.shao@amd.com>
Co-authored-by: Yusheng (Ethan) Su <yushengsu.thu@gmail.com>
Co-authored-by: Hubert Lu <Hubert.Lu@amd.com>
Co-authored-by: xsun <sunxiao04@gmail.com>
This commit is contained in:
Yuzhen Zhou
2025-12-19 02:38:46 -05:00
committed by GitHub
parent f2d64e6782
commit 4bf06635fc
34 changed files with 823 additions and 72 deletions

View File

@@ -366,10 +366,28 @@ def maybe_download_model(
Local path to the model
"""
# If the path exists locally, return it
def _verify_model_complete(path: str) -> bool:
"""Check if model directory has required subdirectories."""
transformer_dir = os.path.join(path, "transformer")
vae_dir = os.path.join(path, "vae")
config_path = os.path.join(path, "model_index.json")
return (
os.path.exists(config_path)
and os.path.exists(transformer_dir)
and os.path.exists(vae_dir)
)
# If the path exists locally, verify it's complete
if os.path.exists(model_name_or_path):
logger.info("Model already exists locally")
return model_name_or_path
if _verify_model_complete(model_name_or_path):
logger.info("Model already exists locally and is complete")
return model_name_or_path
else:
logger.warning(
"Local model at %s appears incomplete (missing transformer/ or vae/), "
"will attempt re-download",
model_name_or_path,
)
# Otherwise, assume it's a HF Hub model ID and try to download it
try:
@@ -385,7 +403,24 @@ def maybe_download_model(
ignore_patterns=["*.onnx", "*.msgpack"],
local_dir=local_dir,
)
logger.info("Downloaded model to %s", local_path)
# Verify downloaded model is complete
if not _verify_model_complete(local_path):
logger.warning(
"Downloaded model at %s is incomplete, retrying with force_download=True",
local_path,
)
with (
get_lock(model_name_or_path).acquire(poll_interval=2),
suppress_other_loggers(not_suppress_on_main_rank=True),
):
local_path = snapshot_download(
repo_id=model_name_or_path,
ignore_patterns=["*.onnx", "*.msgpack"],
local_dir=local_dir,
force_download=True,
)
logger.info("Downloaded model to %s", local_path)
return str(local_path)
except Exception as e:
raise ValueError(