mirror of
https://github.com/huggingface/diffusers.git
synced 2026-10-24 20:59:27 +03:00
acbabca385
* introduce TorchDeviceBackend * update --------- Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
78 lines
2.9 KiB
Python
78 lines
2.9 KiB
Python
# coding=utf-8
|
|
# Copyright 2026 HuggingFace Inc.
|
|
#
|
|
# 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.
|
|
|
|
import json
|
|
import os
|
|
from unittest import mock
|
|
|
|
import pytest
|
|
import torch
|
|
from huggingface_hub import hf_hub_download
|
|
|
|
from diffusers.utils.torch_utils import TorchDeviceBackend
|
|
|
|
|
|
def backend_memory_allocated(device: str) -> int:
|
|
"""
|
|
Bytes currently allocated on `device`. `tests/testing_utils.py` only exposes the *peak* allocation, which cannot
|
|
show memory being released. Skips on backends that do not implement `memory_allocated()` (e.g. mps).
|
|
"""
|
|
device_module = getattr(torch, torch.device(device).type)
|
|
if not hasattr(device_module, "memory_allocated"):
|
|
pytest.skip(f"`memory_allocated()` is not implemented for {device}.")
|
|
return device_module.memory_allocated()
|
|
|
|
|
|
def patch_free_memory(free_bytes: int, total_bytes: int = 80 * 1024):
|
|
"""
|
|
Simulate `free_bytes` of free device memory on whichever backend backs `torch_device`.
|
|
|
|
`TorchDeviceBackend.mem_get_info` returns `(free, total)` and is the single point where `AutoOffloadStrategy` learns how
|
|
much memory is available, so patching it makes offloading decisions deterministic instead of dependent on the real
|
|
free memory of the test hardware (an 80GB GPU never runs low on a handful of KB-sized models).
|
|
"""
|
|
return mock.patch.object(TorchDeviceBackend, "mem_get_info", return_value=(free_bytes, total_bytes))
|
|
|
|
|
|
def get_specified_components(path_or_repo_id, cache_dir=None):
|
|
"""
|
|
The component names a `modular_model_index.json` actually points at a checkpoint for. Returns `None` when the
|
|
index cannot be fetched, which callers treat as "skip the comparison".
|
|
"""
|
|
if os.path.isdir(path_or_repo_id):
|
|
config_path = os.path.join(path_or_repo_id, "modular_model_index.json")
|
|
else:
|
|
try:
|
|
config_path = hf_hub_download(
|
|
repo_id=path_or_repo_id,
|
|
filename="modular_model_index.json",
|
|
local_dir=cache_dir,
|
|
)
|
|
except Exception:
|
|
return None
|
|
|
|
with open(config_path) as f:
|
|
config = json.load(f)
|
|
|
|
components = set()
|
|
for k, v in config.items():
|
|
if isinstance(v, (str, int, float, bool)):
|
|
continue
|
|
for entry in v:
|
|
if isinstance(entry, dict) and (entry.get("repo") or entry.get("pretrained_model_name_or_path")):
|
|
components.add(k)
|
|
break
|
|
return components
|