1
0
mirror of https://github.com/huggingface/diffusers.git synced 2026-10-23 09:09:01 +03:00
Files
diffusers/tests/modular_pipelines/testing_utils/utils.py
Akshan Krithick 62cde2b011 add modular pipeline for qwen-image 2.1 (#14976)
* add modular pipeline for qwen-image 2.1

* split out a turbo blockset and drop the guider from the text encoders

* Update tests/modular_pipelines/qwenimage21/test_modular_pipeline_qwenimage21_turbo.py

* Update tests/modular_pipelines/qwenimage21/test_modular_pipeline_qwenimage21.py

---------

Co-authored-by: Dhruv Nair <dhruv.nair@gmail.com>
2026-10-15 17:52:16 +05:30

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 v is None or 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