import os

from . import custom, diffusers, timm, torchaudio, torchvision, transformers
from .executor import run_fwd, run_fwd_bwd
from .registry import model_zoo

# We pick a subset of models for fast testing in order to reduce the total testing time
COMMON_MODELS = [
    "custom_hanging_param_model",
    "custom_nested_model",
    "custom_repeated_computed_layers",
    "custom_simple_net",
    "diffusers_clip_text_model",
    "diffusers_auto_encoder_kl",
    "diffusers_unet2d_model",
    "timm_densenet",
    "timm_resnet",
    "timm_swin_transformer",
    "torchaudio_wav2vec2_base",
    "torchaudio_conformer",
    "transformers_bert_for_masked_lm",
    "transformers_bloom_for_causal_lm",
    "transformers_falcon_for_causal_lm",
    "transformers_chatglm_for_conditional_generation",
    "transformers_llama_for_causal_lm",
    "transformers_vit_for_masked_image_modeling",
    "transformers_mistral_for_causal_lm",
]

IS_FAST_TEST = os.environ.get("FAST_TEST", "0") == "1"


__all__ = ["model_zoo", "run_fwd", "run_fwd_bwd", "COMMON_MODELS", "IS_FAST_TEST"]