faster test

open-mmlab · Sep 15, 2023 · 6e4ef2a · 6e4ef2a
1 parent b9d39bf
commit 6e4ef2a
Show file tree

Hide file tree

Showing 4 changed files with 21 additions and 21 deletions.
diff --git a/tests/test_models/test_editors/test_controlnet/test_controlnet.py b/tests/test_models/test_editors/test_controlnet/test_controlnet.py
@@ -22,7 +22,7 @@
 
 register_all_modules()
 
-stable_diffusion_v15_url = 'runwayml/stable-diffusion-v1-5'
+stable_diffusion_tiny_url = 'hf-internal-testing/tiny-stable-diffusion-pipe'
 config = dict(
     type='ControlStableDiffusion',
     vae=dict(type='AutoencoderKL', sample_size=64),
@@ -37,21 +37,21 @@
     text_encoder=dict(
         type='ClipWrapper',
         clip_type='huggingface',
-        pretrained_model_name_or_path=stable_diffusion_v15_url,
+        pretrained_model_name_or_path=stable_diffusion_tiny_url,
         subfolder='text_encoder'),
-    tokenizer=stable_diffusion_v15_url,
+    tokenizer=stable_diffusion_tiny_url,
     controlnet=dict(
         type='ControlNetModel',
         # from_pretrained=controlnet_canny_rul
         from_config=config_path  # train from scratch
     ),
     scheduler=dict(
         type='DDPMScheduler',
-        from_pretrained=stable_diffusion_v15_url,
+        from_pretrained=stable_diffusion_tiny_url,
         subfolder='scheduler'),
     test_scheduler=dict(
         type='DDIMScheduler',
-        from_pretrained=stable_diffusion_v15_url,
+        from_pretrained=stable_diffusion_tiny_url,
         subfolder='scheduler'),
     data_preprocessor=dict(type='DataPreprocessor'),
     enable_xformers=False,

diff --git a/tests/test_models/test_editors/test_dreambooth/test_dreambooth.py b/tests/test_models/test_editors/test_dreambooth/test_dreambooth.py
@@ -21,7 +21,7 @@
 
 register_all_modules()
 
-stable_diffusion_v15_url = 'runwayml/stable-diffusion-v1-5'
+stable_diffusion_tiny_url = 'diffusers/tiny-stable-diffusion-torch'
 val_prompts = ['a sks dog in basket']
 
 config = dict(
@@ -38,16 +38,16 @@
     text_encoder=dict(
         type='ClipWrapper',
         clip_type='huggingface',
-        pretrained_model_name_or_path=stable_diffusion_v15_url,
+        pretrained_model_name_or_path=stable_diffusion_tiny_url,
         subfolder='text_encoder'),
-    tokenizer=stable_diffusion_v15_url,
+    tokenizer=stable_diffusion_tiny_url,
     scheduler=dict(
         type='DDPMScheduler',
-        from_pretrained=stable_diffusion_v15_url,
+        from_pretrained=stable_diffusion_tiny_url,
         subfolder='scheduler'),
     test_scheduler=dict(
         type='DDIMScheduler',
-        from_pretrained=stable_diffusion_v15_url,
+        from_pretrained=stable_diffusion_tiny_url,
         subfolder='scheduler'),
     dtype='fp32',
     data_preprocessor=dict(type='DataPreprocessor'),

diff --git a/tests/test_models/test_editors/test_fastcomposer/test_fastcomposer.py b/tests/test_models/test_editors/test_fastcomposer/test_fastcomposer.py
@@ -14,7 +14,7 @@
 gc.collect()
 torch.cuda.empty_cache()
 register_all_modules()
-stable_diffusion_v15_url = 'runwayml/stable-diffusion-v1-5'
+stable_diffusion_tiny_url = 'diffusers/tiny-stable-diffusion-torch'
 finetuned_model_path = ''
 
 vision_config = dict(
@@ -49,13 +49,13 @@
     text_encoder=dict(
         type='ClipWrapper',
         clip_type='huggingface',
-        pretrained_model_name_or_path=stable_diffusion_v15_url,
+        pretrained_model_name_or_path=stable_diffusion_tiny_url,
         subfolder='text_encoder'),
-    tokenizer=stable_diffusion_v15_url,
+    tokenizer=stable_diffusion_tiny_url,
     pretrained_cfg=dict(
         finetuned_model_path=finetuned_model_path,
         enable_xformers_memory_efficient_attention=None,
-        pretrained_model_name_or_path=stable_diffusion_v15_url,
+        pretrained_model_name_or_path=stable_diffusion_tiny_url,
         image_encoder=vision_config,
         revision=None,
         non_ema_revision=None,
@@ -70,11 +70,11 @@
         object_resolution=64),
     scheduler=dict(
         type='DDPMScheduler',
-        from_pretrained=stable_diffusion_v15_url,
+        from_pretrained=stable_diffusion_tiny_url,
         subfolder='scheduler'),
     test_scheduler=dict(
         type='DDIMScheduler',
-        from_pretrained=stable_diffusion_v15_url,
+        from_pretrained=stable_diffusion_tiny_url,
         subfolder='scheduler'),
     dtype='fp32',
     data_preprocessor=dict(type='DataPreprocessor'))

diff --git a/tests/test_models/test_editors/test_vico/test_vico.py b/tests/test_models/test_editors/test_vico/test_vico.py
@@ -21,7 +21,7 @@
 
 register_all_modules()
 
-stable_diffusion_v15_url = 'runwayml/stable-diffusion-v1-5'
+stable_diffusion_tiny_url = 'diffusers/tiny-stable-diffusion-torch'
 val_prompts = ['a photo of S*']
 image_cross_layers = [
     # down blocks (2x transformer block) * (3x down blocks) = 6
@@ -63,16 +63,16 @@
     text_encoder=dict(
         type='ClipWrapper',
         clip_type='huggingface',
-        pretrained_model_name_or_path=stable_diffusion_v15_url,
+        pretrained_model_name_or_path=stable_diffusion_tiny_url,
         subfolder='text_encoder'),
-    tokenizer=stable_diffusion_v15_url,
+    tokenizer=stable_diffusion_tiny_url,
     scheduler=dict(
         type='DDPMScheduler',
-        from_pretrained=stable_diffusion_v15_url,
+        from_pretrained=stable_diffusion_tiny_url,
         subfolder='scheduler'),
     test_scheduler=dict(
         type='DDIMScheduler',
-        from_pretrained=stable_diffusion_v15_url,
+        from_pretrained=stable_diffusion_tiny_url,
         subfolder='scheduler'),
     dtype='fp32',
     data_preprocessor=dict(type='DataPreprocessor'),