Spaces:
Running
on
Zero
Running
on
Zero
| # Copyright 2023 The HuggingFace Team. All rights reserved. | |
| # | |
| # 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. | |
| """ | |
| PyTorch utilities: Utilities related to PyTorch | |
| """ | |
| from typing import List, Optional, Tuple, Union | |
| from . import logging | |
| from .import_utils import is_torch_available, is_torch_version | |
| if is_torch_available(): | |
| import torch | |
| logger = logging.get_logger(__name__) # pylint: disable=invalid-name | |
| try: | |
| from torch._dynamo import allow_in_graph as maybe_allow_in_graph | |
| except (ImportError, ModuleNotFoundError): | |
| def maybe_allow_in_graph(cls): | |
| return cls | |
| def randn_tensor( | |
| shape: Union[Tuple, List], | |
| generator: Optional[Union[List["torch.Generator"], "torch.Generator"]] = None, | |
| device: Optional["torch.device"] = None, | |
| dtype: Optional["torch.dtype"] = None, | |
| layout: Optional["torch.layout"] = None, | |
| ): | |
| """A helper function to create random tensors on the desired `device` with the desired `dtype`. When | |
| passing a list of generators, you can seed each batch size individually. If CPU generators are passed, the tensor | |
| is always created on the CPU. | |
| """ | |
| # device on which tensor is created defaults to device | |
| rand_device = device | |
| batch_size = shape[0] | |
| layout = layout or torch.strided | |
| device = device or torch.device("cpu") | |
| if generator is not None: | |
| gen_device_type = generator.device.type if not isinstance(generator, list) else generator[0].device.type | |
| if gen_device_type != device.type and gen_device_type == "cpu": | |
| rand_device = "cpu" | |
| if device != "mps": | |
| logger.info( | |
| f"The passed generator was created on 'cpu' even though a tensor on {device} was expected." | |
| f" Tensors will be created on 'cpu' and then moved to {device}. Note that one can probably" | |
| f" slighly speed up this function by passing a generator that was created on the {device} device." | |
| ) | |
| elif gen_device_type != device.type and gen_device_type == "cuda": | |
| raise ValueError(f"Cannot generate a {device} tensor from a generator of type {gen_device_type}.") | |
| if isinstance(generator, list): | |
| shape = (1,) + shape[1:] | |
| latents = [ | |
| torch.randn(shape, generator=generator[i], device=rand_device, dtype=dtype, layout=layout) | |
| for i in range(batch_size) | |
| ] | |
| latents = torch.cat(latents, dim=0).to(device) | |
| else: | |
| latents = torch.randn(shape, generator=generator, device=rand_device, dtype=dtype, layout=layout).to(device) | |
| return latents | |
| def is_compiled_module(module): | |
| """Check whether the module was compiled with torch.compile()""" | |
| if is_torch_version("<", "2.0.0") or not hasattr(torch, "_dynamo"): | |
| return False | |
| return isinstance(module, torch._dynamo.eval_frame.OptimizedModule) | |