github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-11 13:59:08 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Jiarui Fang
2022-11-08 11:45:23 +08:00
committed by GitHub
parent 203ca57aed
commit 6fa71d65d3
2 changed files with 9 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
tests/test_fx/test_tracer/test_hf_model/test_hf_diffuser.py
View File

@@ -1,12 +1,17 @@
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 = [