Unverified Commit 6fa71d65 authored by Jiarui Fang's avatar Jiarui Fang Committed by GitHub
Browse files

[fx] skip diffusers unitest if it is not installed (#1799)

parent 203ca57a
diffusers
fbgemm-gpu==0.2.0
pytest
torchvision
......
import diffusers
import pytest
import torch
import transformers
from torch.fx import GraphModule
from utils import trace_model_and_compare_output
import transformers
from colossalai.fx import ColoTracer
try:
import diffusers
HAS_DIFFUSERS = True
except ImportError:
HAS_DIFFUSERS = False
BATCH_SIZE = 2
SEQ_LENGTH = 5
HEIGHT = 224
......@@ -16,6 +21,7 @@ LATENTS_SHAPE = (BATCH_SIZE, IN_CHANNELS, HEIGHT // 8, WIDTH // 8)
TIME_STEP = 2
@pytest.mark.skipif(not HAS_DIFFUSERS, reason="diffusers has not been installed")
def test_vae():
MODEL_LIST = [
diffusers.AutoencoderKL,
......@@ -80,6 +86,7 @@ def test_clip():
trace_model_and_compare_output(model, data_gen)
@pytest.mark.skipif(not HAS_DIFFUSERS, reason="diffusers has not been installed")
@pytest.mark.skip(reason='cannot pass the test yet')
def test_unet():
MODEL_LIST = [
......
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