StableCascade image_encoder

Author: saddam213

OnnxStack.Converter/stable_cascade/config_image_encoder.json

@@ -0,0 +1,113 @@
+{
+    "input_model": {
+        "type": "PyTorchModel",
+        "config": {
+            "model_path": "stabilityai/stable-cascade",
+            "model_loader": "image_encoder_load",
+            "model_script": "models.py",
+            "io_config": {
+                "input_names": [ "sample"],
+                "output_names": [ "latent_sample" ],
+                "dynamic_axes": { "sample": { "0": "batch", "1": "channels", "2": "height", "3": "width" } }
+            },
+            "dummy_inputs_func": "image_encoder_conversion_inputs"
+        }
+    },
+    "systems": {
+        "local_system": {
+            "type": "LocalSystem",
+            "config": {
+                "accelerators": [
+                    {
+                        "device": "gpu",
+                        "execution_providers": [
+                            "DmlExecutionProvider"
+                        ]
+                    }
+                ]
+            }
+        }
+    },
+    "evaluators": {
+        "common_evaluator": {
+            "metrics": [
+                {
+                    "name": "latency",
+                    "type": "latency",
+                    "sub_types": [{"name": "avg"}],
+                    "user_config": {
+                        "user_script": "models.py",
+                        "dataloader_func": "image_encoder_data_loader",
+                        "batch_size": 1
+                    }
+                }
+            ]
+        }
+    },
+    "passes": {
+        "convert": {
+            "type": "OnnxConversion",
+            "config": {
+                "target_opset": 16
+            }
+        },
+        "optimize": {
+            "type": "OrtTransformersOptimization",
+            "config": {
+                "model_type": "clip",
+                "opt_level": 0,
+                "float16": true,
+                "use_gpu": true,
+                "keep_io_types": true,
+                "optimization_options": {
+                    "enable_gelu": true,
+                    "enable_layer_norm": true,
+                    "enable_attention": true,
+                    "use_multi_head_attention": true,
+                    "enable_skip_layer_norm": false,
+                    "enable_embed_layer_norm": true,
+                    "enable_bias_skip_layer_norm": false,
+                    "enable_bias_gelu": true,
+                    "enable_gelu_approximation": false,
+                    "enable_qordered_matmul": false,
+                    "enable_shape_inference": true,
+                    "enable_gemm_fast_gelu": false,
+                    "enable_nhwc_conv": false,
+                    "enable_group_norm": true,
+                    "enable_bias_splitgelu": false,
+                    "enable_packed_qkv": true,
+                    "enable_packed_kv": true,
+                    "enable_bias_add": false,
+                    "group_norm_channels_last": false
+                },
+                "force_fp32_ops": ["RandomNormalLike"],
+                "force_fp16_inputs": {
+                    "GroupNorm": [0, 1, 2]
+                }
+            }
+        },
+        "optimize_cuda": {
+            "type": "OrtTransformersOptimization",
+            "config": {
+                "model_type": "clip",
+                "opt_level": 0,
+                "float16": true,
+                "use_gpu": true,
+                "keep_io_types": false
+            }
+        }
+    },
+    "pass_flows": [
+        ["convert", "optimize"]
+    ],
+    "engine": {
+        "log_severity_level": 0,
+        "evaluator": "common_evaluator",
+        "evaluate_input_model": false,
+        "host": "local_system",
+        "target": "local_system",
+        "cache_dir": "cache",
+        "output_name": "image_encoder",
+        "output_dir": "footprints"
+    }
+}

OnnxStack.Converter/stable_cascade/convert.py

@@ -121,7 +121,7 @@ def optimize(
 
     model_info = {}
 
-    submodel_names = [ "text_encoder", "decoder", "prior"]
+    submodel_names = [ "text_encoder", "decoder", "prior", "image_encoder"]
 
     has_safety_checker = getattr(pipeline, "safety_checker", None) is not None
 

OnnxStack.Converter/stable_cascade/models.py

@@ -6,7 +6,7 @@
 import torch
 from typing import Union, Optional, Tuple
 from diffusers import AutoencoderKL, StableCascadeUNet
-from transformers.models.clip.modeling_clip import CLIPTextModelWithProjection
+from transformers.models.clip.modeling_clip import CLIPTextModelWithProjection, CLIPVisionModelWithProjection
 from dataclasses import dataclass
 
 # Helper latency-only dataloader that creates random tensors with no label
@@ -111,4 +111,31 @@ def prior_conversion_inputs(model=None):
 
 
 def prior_data_loader(data_dir, batchsize, *args, **kwargs):
-    return RandomDataLoader(prior_inputs, batchsize, torch.float16)
+    return RandomDataLoader(prior_inputs, batchsize, torch.float16)
+
+
+
+
+    
+# -----------------------------------------------------------------------------
+# image_encoder
+# -----------------------------------------------------------------------------
+
+def image_encoder_inputs(batchsize, torch_dtype, is_conversion_inputs=False):
+    inputs = {
+        "sample": torch.rand((batchsize, 3, 224, 224), dtype=torch_dtype)
+    }
+    return inputs
+
+
+def image_encoder_load(model_name):
+    model = CLIPVisionModelWithProjection.from_pretrained(model_name, subfolder="image_encoder")
+    return model
+
+
+def image_encoder_conversion_inputs(model=None):
+    return tuple(image_encoder_inputs(1, torch.float32, True).values())
+
+
+def image_encoder_data_loader(data_dir, batchsize, *args, **kwargs):
+    return RandomDataLoader(image_encoder_inputs, batchsize, torch.float16)