StableCascade vqgan

Author: saddam213

OnnxStack.Converter/stable_cascade/config_vqgan.json

@@ -0,0 +1,103 @@
+{
+    "input_model": {
+        "type": "PyTorchModel",
+        "config": {
+            "model_path": "stabilityai/stable-cascade",
+            "model_loader": "vqgan_load",
+            "model_script": "models.py",
+            "io_config": {
+                "input_names": [ "sample", "return_dict" ],
+                "output_names": [ "latent_sample" ],
+                "dynamic_axes": { "sample": { "0": "batch", "1": "channels", "2": "height", "3": "width" } }
+            },
+            "dummy_inputs_func": "vqgan_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": "vqgan_data_loader",
+                        "batch_size": 1
+                    }
+                }
+            ]
+        }
+    },
+    "passes": {
+        "convert": {
+            "type": "OnnxConversion",
+            "config": {
+                "target_opset": 16
+            }
+        },
+        "optimize": {
+            "type": "OrtTransformersOptimization",
+            "config": {
+                "model_type": "vae",
+                "opt_level": 0,
+                "float16": true,
+                "use_gpu": true,
+                "keep_io_types": false,
+                "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]
+                }
+            }
+        }
+    },
+    "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": "vqgan",
+        "output_dir": "footprints"
+    }
+}

OnnxStack.Converter/stable_cascade/models.py

@@ -6,6 +6,7 @@
 import torch
 from typing import Union, Optional, Tuple
 from diffusers import AutoencoderKL, StableCascadeUNet
+from diffusers.pipelines.wuerstchen import PaellaVQModel
 from transformers.models.clip.modeling_clip import CLIPTextModelWithProjection, CLIPVisionModelWithProjection
 from dataclasses import dataclass
 
@@ -129,7 +130,7 @@ def image_encoder_inputs(batchsize, torch_dtype, is_conversion_inputs=False):
 
 
 def image_encoder_load(model_name):
-    model = CLIPVisionModelWithProjection.from_pretrained(model_name, subfolder="image_encoder")
+    model = CLIPVisionModelWithProjection.from_pretrained(model_name, subfolder="image_encoder", use_safetensors=True)
     return model
 
 
@@ -138,4 +139,28 @@ def image_encoder_conversion_inputs(model=None):
 
 
 def image_encoder_data_loader(data_dir, batchsize, *args, **kwargs):
-    return RandomDataLoader(image_encoder_inputs, batchsize, torch.float16)
+    return RandomDataLoader(image_encoder_inputs, batchsize, torch.float16)
+
+
+# -----------------------------------------------------------------------------
+# vqgan
+# -----------------------------------------------------------------------------
+
+def vqgan_inputs(batchsize, torch_dtype, is_conversion_inputs=False):
+    inputs = {
+        "sample": torch.rand((batchsize, 3, 256, 256), dtype=torch_dtype)
+    }
+    return inputs
+
+
+def vqgan_load(model_name):
+    model = PaellaVQModel.from_pretrained(model_name, subfolder="vqgan", use_safetensors=True)
+    return model
+
+
+def vqgan_conversion_inputs(model=None):
+    return tuple(vqgan_inputs(1, torch.float32, True).values())
+
+
+def vqgan_data_loader(data_dir, batchsize, *args, **kwargs):
+    return RandomDataLoader(vqgan_inputs, batchsize, torch.float16)