[tests] diffuser models in model zoo (#3136)

* [tests] diffuser models in model zoo

* remove useless code

* [tests] add diffusers to requirement-test

tests/kit/model_zoo/diffusers/__init__.py

@@ -0,0 +1 @@
+from .diffusers import *

tests/kit/model_zoo/diffusers/diffusers.py

@@ -0,0 +1,73 @@
+from functools import partial
+
+import diffusers
+import torch
+import transformers
+
+from ..registry import ModelAttribute, model_zoo
+
+BATCH_SIZE = 2
+SEQ_LENGTH = 5
+HEIGHT = 224
+WIDTH = 224
+IN_CHANNELS = 3
+LATENTS_SHAPE = (BATCH_SIZE, IN_CHANNELS, HEIGHT // 7, WIDTH // 7)
+TIME_STEP = 3
+
+data_vae_fn = lambda: dict(sample=torch.randn(2, 3, 32, 32))
+data_unet_fn = lambda: dict(sample=torch.randn(2, 3, 32, 32), timestep=3)
+
+identity_output = lambda x: x
+
+
+def data_clip_model():
+    input_ids = torch.zeros((BATCH_SIZE, SEQ_LENGTH), dtype=torch.int64)
+    attention_mask = torch.zeros((BATCH_SIZE, SEQ_LENGTH), dtype=torch.int64)
+    position_ids = torch.zeros((BATCH_SIZE, SEQ_LENGTH), dtype=torch.int64)
+    pixel_values = torch.zeros((BATCH_SIZE, IN_CHANNELS, HEIGHT, WIDTH), dtype=torch.float32)
+    return dict(input_ids=input_ids,
+                pixel_values=pixel_values,
+                attention_mask=attention_mask,
+                position_ids=position_ids)
+
+
+def data_clip_text():
+    input_ids = torch.zeros((BATCH_SIZE, SEQ_LENGTH), dtype=torch.int64)
+    attention_mask = torch.zeros((BATCH_SIZE, SEQ_LENGTH), dtype=torch.int64)
+    return dict(input_ids=input_ids, attention_mask=attention_mask)
+
+
+def data_clip_vision():
+    pixel_values = torch.zeros((BATCH_SIZE, IN_CHANNELS, HEIGHT, WIDTH), dtype=torch.float32)
+    return dict(pixel_values=pixel_values)
+
+
+model_zoo.register(name='diffusers_auto_encoder_kl',
+                   model_fn=diffusers.AutoencoderKL,
+                   data_gen_fn=data_vae_fn,
+                   output_transform_fn=identity_output)
+
+model_zoo.register(name='diffusers_vq_model',
+                   model_fn=diffusers.VQModel,
+                   data_gen_fn=data_vae_fn,
+                   output_transform_fn=identity_output)
+
+model_zoo.register(name='diffusers_clip_model',
+                   model_fn=partial(transformers.CLIPModel, config=transformers.CLIPConfig()),
+                   data_gen_fn=data_clip_model,
+                   output_transform_fn=identity_output)
+
+model_zoo.register(name='diffusers_clip_text_model',
+                   model_fn=partial(transformers.CLIPTextModel, config=transformers.CLIPTextConfig()),
+                   data_gen_fn=data_clip_text,
+                   output_transform_fn=identity_output)
+
+model_zoo.register(name='diffusers_clip_vision_model',
+                   model_fn=partial(transformers.CLIPVisionModel, config=transformers.CLIPVisionConfig()),
+                   data_gen_fn=data_clip_vision,
+                   output_transform_fn=identity_output)
+
+model_zoo.register(name='diffusers_unet2d_model',
+                   model_fn=diffusers.UNet2DModel,
+                   data_gen_fn=data_unet_fn,
+                   output_transform_fn=identity_output)