1
0
mirror of https://github.com/huggingface/diffusers.git synced 2026-01-29 07:22:12 +03:00

fix torchvision import (#6796)

This commit is contained in:
Patrick von Platen
2024-02-01 00:13:10 +02:00
committed by GitHub
parent 04cd6adf8c
commit c3369f5673

View File

@@ -5,7 +5,7 @@ from typing import Any, Dict, Iterable, List, Optional, Union
import numpy as np
import torch
from torchvision import transforms
from transformers import is_torchvision_available
from .models import UNet2DConditionModel
from .utils import (
@@ -23,6 +23,9 @@ if is_transformers_available():
if is_peft_available():
from peft import set_peft_model_state_dict
if is_torchvision_available():
from torchvision import transforms
def set_seed(seed: int):
"""
@@ -79,6 +82,11 @@ def resolve_interpolation_mode(interpolation_type: str):
`torchvision.transforms.InterpolationMode`: an `InterpolationMode` enum used by torchvision's `resize`
transform.
"""
if not is_torchvision_available():
raise ImportError(
"Please make sure to install `torchvision` to be able to use the `resolve_interpolation_mode()` function."
)
if interpolation_type == "bilinear":
interpolation_mode = transforms.InterpolationMode.BILINEAR
elif interpolation_type == "bicubic":