Unverified Commit c3369f56 authored by Patrick von Platen's avatar Patrick von Platen Committed by GitHub
Browse files

fix torchvision import (#6796)

parent 04cd6adf
......@@ -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":
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment