Initial commit with folder contents

.gitattributes

@@ -33,8 +33,329 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
-cached_pipe/text_encoder filter=lfs diff=lfs merge=lfs -text
-cached_pipe/text_encoder_2 filter=lfs diff=lfs merge=lfs -text
-cached_pipe/vae.decoder filter=lfs diff=lfs merge=lfs -text
-cached_pipe/fast_unet filter=lfs diff=lfs merge=lfs -text
-cached_pipe/unet filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_6936 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_6920 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_6912 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7600 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_6955 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_6958 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/resnets.0.time_emb_proj.weight filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7203 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7624 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_6910 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_6969 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_6909 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_6911 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_6935 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7003 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_6979 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/resnets.2.time_emb_proj.weight filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7193 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_6921 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7017 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7609 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_6982 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7169 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_6984 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7182 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7031 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7041 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_6934 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7619 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_6931 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_6968 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7160 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7008 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_6945 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7632 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7168 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7007 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7030 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/resnets.1.time_emb_proj.weight filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7103 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7040 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7065 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7157 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7055 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7006 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7089 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_6993 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/resnets.2.conv_shortcut.weight filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7054 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7099 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7158 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7137 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7027 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7147 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7112 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7128 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7104 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7016 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_6960 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_6994 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7126 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7018 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_6922 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_6970 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7611 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7042 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_6947 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7019 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7634 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7192 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_6971 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/resnets.0.conv_shortcut.weight filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7633 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7090 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7241 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7032 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_6995 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_6946 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7227 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7635 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7232 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7043 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7136 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7587 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7127 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7088 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7254 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_6923 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7080 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7056 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7123 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7075 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/resnets.1.conv_shortcut.weight filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7064 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7051 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7079 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7347 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7371 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7102 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7078 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7336 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7360 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7361 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7179 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7159 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7184 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7113 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7328 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7376 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7337 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7115 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7171 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7407 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7256 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7216 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7326 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7251 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7208 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7303 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7217 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7350 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7231 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7183 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7265 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7240 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7206 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7323 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7207 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7289 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7091 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7230 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7066 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_6944 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7067 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7290 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_6983 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7139 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7302 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_6992 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7315 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_6959 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7299 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7386 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7138 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7363 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7312 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7114 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7219 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7150 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7195 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7351 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7264 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7327 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7280 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7170 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7313 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7255 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7242 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7278 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7291 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7279 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7288 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7304 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7352 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7218 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7499 filter=lfs diff=lfs merge=lfs -text
+engine/down_blocks.0.plan filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7432 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7503 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7480 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7488 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7464 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7475 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7502 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7513 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7527 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7537 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7275 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7430 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7267 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7465 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7405 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7512 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7339 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7427 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7194 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7243 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7621 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7266 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7456 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7395 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7077 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7417 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7585 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7314 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7623 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7443 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7419 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7643 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7489 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7455 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7479 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7549 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_6957 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_6933 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7504 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7431 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7441 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7440 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7408 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7466 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7416 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7029 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7406 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7552 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7451 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7053 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7005 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7646 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7560 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7338 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7595 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7362 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7574 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7454 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7645 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7622 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7478 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7598 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7608 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7575 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7561 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7563 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7375 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7385 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7584 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7384 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7467 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7599 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7551 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7571 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7125 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7442 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7528 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7576 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7398 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7418 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7562 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7374 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7610 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7101 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7523 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7526 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7536 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7586 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_6981 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7387 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7277 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7547 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7539 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7550 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7515 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7491 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7514 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7538 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7149 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7490 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7229 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7573 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7597 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7181 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7373 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7253 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7349 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7301 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7325 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7205 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7501 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7429 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7397 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7453 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_6932 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7348 filter=lfs diff=lfs merge=lfs -text
+engine/up_blocks.2.plan filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7525 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7477 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7500 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7644 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_6980 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7428 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_6956 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7572 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7620 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/resnets.0.conv2.weight filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7004 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/resnets.2.conv2.weight filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7052 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7148 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7252 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7100 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7300 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7452 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/upsamplers.0.conv.weight filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7524 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7476 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7028 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/resnets.1.conv2.weight filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7124 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7396 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7076 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7372 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7180 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7228 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7204 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7548 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/resnets.2.conv1.weight filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7324 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7276 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/onnx__MatMul_7596 filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/resnets.1.conv1.weight filter=lfs diff=lfs merge=lfs -text
+onnx/up_blocks.0/resnets.0.conv1.weight filter=lfs diff=lfs merge=lfs -text
+engine/down_blocks.1.plan filter=lfs diff=lfs merge=lfs -text
+engine/up_blocks.1.plan filter=lfs diff=lfs merge=lfs -text
+engine/mid_block.plan filter=lfs diff=lfs merge=lfs -text
+engine/down_blocks.2.plan filter=lfs diff=lfs merge=lfs -text
+engine/up_blocks.0.plan filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,2 @@
+**/__pycache__
+**.egg-info

.gitmodules

@@ -0,0 +1,3 @@
+[submodule "models/newdream-sdxl-20"]
+	path = models/newdream-sdxl-20
+	url = https://huggingface.co/stablediffusionapi/newdream-sdxl-20

engine/down_blocks.0.plan

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5c4b70d1a0416aa494a3d7759349f9637c36c4e26729c3190f5a16bfc497694b
+size 11712396

engine/down_blocks.1.plan

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:21baa4ea0b4740f3cec8c35825dd6cc4dca7dc39cda3f6be8eab2ce1cd3e834f
+size 124421828

engine/down_blocks.2.plan

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4412bf1d85e21eaac612aa9319b33f35b7e83d7858cca597bb6e645fa7bc7207
+size 1522617884

engine/mid_block.plan

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:040b2e3dddac9efd68e125b02056a5cd4f8df7259e7167f14e57db8b68e26fb8
+size 830401652

engine/up_blocks.0.plan

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0cbc51a0d7f00e4d8b108bfda46a25df62b41623e084733173730774e7e15048
+size 2425023084

engine/up_blocks.1.plan

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1bd8a45dd815e72cdf75dd5df10a651c892b09f80a292edaafd4d192be419be3
+size 218672972

engine/up_blocks.2.plan

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1ee5d0d133198b30624c3d682b0438fa4485eeade489e7ab21ed425e1d947549
+size 24347780

loss_params.pth

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:27ba04dc09bfe8325c2b8d8acbfa5fbf746f61169cf1cdfe07d028ad697217f1
-size 3568
+oid sha256:4e4c687fb455b7495e325d5f1761391d281323de6d2a493b153a3dac9536664e
+size 3120

pyproject.toml

@@ -8,30 +8,21 @@ description = "An edge-maxxing model submission for the 4090 newdream contest"
 requires-python = ">=3.10,<3.11"
 version = "6"
 dependencies = [
-    "diffusers==0.28.2",
-    "onediff==1.2.0",
-    "onediffx==1.2.0",
-    "accelerate==0.31.0",
-    "numpy==1.26.4",
-    "xformers==0.0.25.post1",
-    "triton==2.2.0",
+    "wheel",
+    "diffusers==0.30.2",
     "transformers==4.41.2",
     "accelerate==0.31.0",
     "omegaconf==2.3.0",
-    "torch==2.2.2",
-    "torchvision==0.17.2",
-    "huggingface-hub==0.25.2",
-    "edge-maxxing-pipelines @ git+https://github.com/womboai/edge-maxxing#subdirectory=pipelines",
-    "oneflow",
+    "torch==2.4.1",
+    "edge-maxxing-pipelines @ git+https://github.com/womboai/edge-maxxing@8d8ff45863416484b5b4bc547782591bbdfc696a#subdirectory=pipelines",
+    "polygraphy",
+    "onnx",
+    "tensorrt>=10.5.0",
+    "tensorrt-cu12-libs>=10.5.0",
+    "tensorrt-cu12-bindings>=10.5.0",
+    "cuda-python>=12.6.0",
     "setuptools>=75.2.0",
-    "bitsandbytes>=0.44.1",
-    "stable-fast",
-    "tomesd>=0.1.3",
 ]
 
-[tool.uv.sources]
-oneflow = { url = "https://github.com/siliconflow/oneflow_releases/releases/download/community_cu118/oneflow-0.9.1.dev20240802%2Bcu118-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl" }
-stable-fast = { url = "https://github.com/chengzeyi/stable-fast/releases/download/v1.0.5/stable_fast-1.0.5+torch222cu121-cp310-cp310-manylinux2014_x86_64.whl" }
-
 [project.scripts]
 start_inference = "main:main"

src/cache_diffusion/cachify.py

@@ -0,0 +1,144 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: MIT
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+
+import fnmatch
+from contextlib import contextmanager
+
+from diffusers.models.attention import BasicTransformerBlock, JointTransformerBlock
+from diffusers.models.transformers.pixart_transformer_2d import PixArtTransformer2DModel
+from diffusers.models.transformers.transformer_sd3 import SD3Transformer2DModel
+from diffusers.models.unets.unet_2d_blocks import (
+    CrossAttnDownBlock2D,
+    CrossAttnUpBlock2D,
+    DownBlock2D,
+    UNetMidBlock2DCrossAttn,
+    UpBlock2D,
+)
+from diffusers.models.unets.unet_2d_condition import UNet2DConditionModel
+from diffusers.models.unets.unet_3d_blocks import (
+    CrossAttnDownBlockSpatioTemporal,
+    CrossAttnUpBlockSpatioTemporal,
+    DownBlockSpatioTemporal,
+    UNetMidBlockSpatioTemporal,
+    UpBlockSpatioTemporal,
+)
+from diffusers.models.unets.unet_spatio_temporal_condition import UNetSpatioTemporalConditionModel
+
+from .module import CachedModule
+from .utils import replace_module
+
+CACHED_PIPE = {
+    UNet2DConditionModel: (
+        DownBlock2D,
+        CrossAttnDownBlock2D,
+        UNetMidBlock2DCrossAttn,
+        CrossAttnUpBlock2D,
+        UpBlock2D,
+    ),
+    PixArtTransformer2DModel: (BasicTransformerBlock),
+    UNetSpatioTemporalConditionModel: (
+        CrossAttnDownBlockSpatioTemporal,
+        DownBlockSpatioTemporal,
+        UpBlockSpatioTemporal,
+        CrossAttnUpBlockSpatioTemporal,
+        UNetMidBlockSpatioTemporal,
+    ),
+    SD3Transformer2DModel: (JointTransformerBlock),
+}
+
+
+def _apply_to_modules(model, action, modules=None, config_list=None):
+    if hasattr(model, "use_trt_infer") and model.use_trt_infer:
+        for key, module in model.engines.items():
+            if isinstance(module, CachedModule):
+                action(module)
+            elif config_list:
+                for config in config_list:
+                    if _pass(key, config["wildcard_or_filter_func"]):
+                        model.engines[key] = CachedModule(module, config["select_cache_step_func"])
+    else:
+        for name, module in model.named_modules():
+            if isinstance(module, CachedModule):
+                action(module)
+            elif modules and config_list:
+                for config in config_list:
+                    if _pass(name, config["wildcard_or_filter_func"]) and isinstance(
+                        module, modules
+                    ):
+                        replace_module(
+                            model,
+                            name,
+                            CachedModule(module, config["select_cache_step_func"]),
+                        )
+
+
+def cachify(model, config_list, modules):
+    def cache_action(module):
+        pass  # No action needed, caching is handled in the loop itself
+
+    _apply_to_modules(model, cache_action, modules, config_list)
+
+
+def disable(pipe):
+    model = get_model(pipe)
+    _apply_to_modules(model, lambda module: module.disable_cache())
+
+
+def enable(pipe):
+    model = get_model(pipe)
+    _apply_to_modules(model, lambda module: module.enable_cache())
+
+
+def reset_status(pipe):
+    model = get_model(pipe)
+    _apply_to_modules(model, lambda module: setattr(module, "cur_step", 0))
+
+
+def _pass(name, wildcard_or_filter_func):
+    if isinstance(wildcard_or_filter_func, str):
+        return fnmatch.fnmatch(name, wildcard_or_filter_func)
+    elif callable(wildcard_or_filter_func):
+        return wildcard_or_filter_func(name)
+    else:
+        raise NotImplementedError(f"Unsupported type {type(wildcard_or_filter_func)}")
+
+
+def get_model(pipe):
+    if hasattr(pipe, "unet"):
+        return pipe.unet
+    elif hasattr(pipe, "transformer"):
+        return pipe.transformer
+    else:
+        raise KeyError
+
+
+@contextmanager
+def infer(pipe):
+    try:
+        yield pipe
+    finally:
+        reset_status(pipe)
+
+
+def prepare(pipe, config_list):
+    model = get_model(pipe)
+    assert model.__class__ in CACHED_PIPE.keys(), f"{model.__class__} is not supported!"
+    cachify(model, config_list, CACHED_PIPE[model.__class__])

src/cache_diffusion/module.py

@@ -0,0 +1,55 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: MIT
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+
+from torch import nn
+
+
+class CachedModule(nn.Module):
+    def __init__(self, block, select_cache_step_func) -> None:
+        super().__init__()
+        self.block = block
+        self.select_cache_step_func = select_cache_step_func
+        self.cur_step = 0
+        self.cached_results = None
+        self.enabled = True
+
+    def __getattr__(self, name):
+        try:
+            return super().__getattr__(name)
+        except AttributeError:
+            return getattr(self.block, name)
+
+    def if_cache(self):
+        return self.select_cache_step_func(self.cur_step) and self.enabled
+
+    def enable_cache(self):
+        self.enabled = True
+
+    def disable_cache(self):
+        self.enabled = False
+        self.cur_step = 0
+
+    def forward(self, *args, **kwargs):
+        if not self.if_cache():
+            self.cached_results = self.block(*args, **kwargs)
+        if self.enabled:
+            self.cur_step += 1
+        return self.cached_results

src/cache_diffusion/utils.py

@@ -0,0 +1,61 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: MIT
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+
+import re
+
+SDXL_DEFAULT_CONFIG = [
+    {
+        "wildcard_or_filter_func": lambda name: "up_blocks.2" not in name,
+        "select_cache_step_func": lambda step: (step % 2) != 0,
+    }
+]
+
+PIXART_DEFAULT_CONFIG = [
+    {
+        "wildcard_or_filter_func": lambda name: not re.search(
+            r"transformer_blocks\.(2[1-7])\.", name
+        ),
+        "select_cache_step_func": lambda step: (step % 3) != 0,
+    }
+]
+
+SVD_DEFAULT_CONFIG = [
+    {
+        "wildcard_or_filter_func": lambda name: "up_blocks.3" not in name,
+        "select_cache_step_func": lambda step: (step % 2) != 0,
+    }
+]
+
+SD3_DEFAULT_CONFIG = [
+    {
+        "wildcard_or_filter_func": lambda name: re.search(
+            r"^((?!transformer_blocks\.(1[6-9]|2[0-3])).)*$", name
+        ),
+        "select_cache_step_func": lambda step: (step % 2) != 0,
+    }
+]
+
+
+def replace_module(parent, name_path, new_module):
+    path_parts = name_path.split(".")
+    for part in path_parts[:-1]:
+        parent = getattr(parent, part)
+    setattr(parent, path_parts[-1], new_module)

src/pipeline.py

@@ -1,982 +1,54 @@
 import torch
-from PIL import Image
+from pathlib import Path
+from PIL.Image import Image
+from diffusers import StableDiffusionXLPipeline, DDIMScheduler
 from pipelines.models import TextToImageRequest
 from torch import Generator
-import json
-from diffusers import StableDiffusionXLPipeline, DDIMScheduler
-import inspect
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from cache_diffusion import cachify
+from trt_pipeline.deploy import load_unet_trt
 from loss import SchedulerWrapper
-from onediffx import compile_pipe,load_pipe
-# Import necessary components
-from transformers import (
-    CLIPImageProcessor,
-    CLIPTextModel,
-    CLIPTextModelWithProjection,
-    CLIPTokenizer,
-    CLIPVisionModelWithProjection,
-)
-
-
-from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback
-from diffusers.image_processor import PipelineImageInput, VaeImageProcessor
-from diffusers.loaders import (
-    FromSingleFileMixin,
-    IPAdapterMixin,
-    StableDiffusionXLLoraLoaderMixin,
-    TextualInversionLoaderMixin,
-)
-from diffusers.models import AutoencoderKL, ImageProjection, UNet2DConditionModel
-from diffusers.models.attention_processor import (
-    AttnProcessor2_0,
-    FusedAttnProcessor2_0,
-    XFormersAttnProcessor,
-)
-from diffusers.models.lora import adjust_lora_scale_text_encoder
-from diffusers.schedulers import KarrasDiffusionSchedulers
-from diffusers.utils import (
-    USE_PEFT_BACKEND,
-    deprecate,
-    is_invisible_watermark_available,
-    is_torch_xla_available,
-    logging,
-    replace_example_docstring,
-    scale_lora_layers,
-    unscale_lora_layers,
-)
-from diffusers.utils.torch_utils import randn_tensor
-from diffusers.pipelines.pipeline_utils import DiffusionPipeline, StableDiffusionMixin
-from diffusers.pipelines.stable_diffusion_xl.pipeline_output import StableDiffusionXLPipelineOutput
-
-# Import watermark if available
-if is_invisible_watermark_available():
-    from .watermark import StableDiffusionXLWatermarker
-
-# Check for XLA availability
-if is_torch_xla_available():
-    import torch_xla.core.xla_model as xm
-    XLA_AVAILABLE = True
-else:
-    XLA_AVAILABLE = False
-
-logger = logging.get_logger(__name__)
-
-# Constants
-EXAMPLE_DOC_STRING = """
-    Examples:
-        ```py
-        >>> import torch
-        >>> from diffusers import StableDiffusionXLPipeline
-
-        >>> pipe = StableDiffusionXLPipeline.from_pretrained(
-        >>>     "stabilityai/stable-diffusion-xl-base-1.0", 
-        >>>     torch_dtype=torch.float16
-        >>> )
-        >>> pipe = pipe.to("cuda")
-
-        >>> prompt = "a photo of an astronaut riding a horse on mars"
-        >>> image = pipe(prompt).images[0]
-        ```
-"""
-
-# Helper functions
-def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
-    """Rescale noise configuration."""
-    std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True)
-    std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)
-    noise_pred_rescaled = noise_cfg * (std_text / std_cfg)
-    noise_cfg = guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg
-    return noise_cfg
-
-# Utils functions
-import numpy as np
-def max_pixel_filter(image: Image) -> Image:
-    try:
-        # Convert the image to a numpy array
-        img_array = np.array(image)
-        # Find the maximum pixel value in the image
-        max_val = img_array.max()
-        # Reduce the maximum value to 1
-        img_array[img_array == max_val] -= 2
-        # Convert the numpy array back to an image
-        filtered_image = Image.fromarray(img_array)
-        return filtered_image
-    except:
-        return image
-    
-# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
-def retrieve_timesteps(
-    scheduler,
-    num_inference_steps: Optional[int] = None,
-    device: Optional[Union[str, torch.device]] = None,
-    timesteps: Optional[List[int]] = None,
-    sigmas: Optional[List[float]] = None,
-    **kwargs,
-):
-    if timesteps is not None and sigmas is not None:
-        raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
-    if timesteps is not None:
-        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
-        if not accepts_timesteps:
-            raise ValueError(
-                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
-                f" timestep schedules. Please check whether you are using the correct scheduler."
-            )
-        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
-        timesteps = scheduler.timesteps
-        num_inference_steps = len(timesteps)
-    elif sigmas is not None:
-        accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
-        if not accept_sigmas:
-            raise ValueError(
-                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
-                f" sigmas schedules. Please check whether you are using the correct scheduler."
-            )
-        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
-        timesteps = scheduler.timesteps
-        num_inference_steps = len(timesteps)
-    else:
-        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
-        timesteps = scheduler.timesteps
-    return timesteps, num_inference_steps
-
-
-class StableDiffusionXLPipeline_new(
-    DiffusionPipeline,
-    StableDiffusionMixin,
-    FromSingleFileMixin,
-    StableDiffusionXLLoraLoaderMixin,
-    TextualInversionLoaderMixin,
-    IPAdapterMixin,
-):
-
-    model_cpu_offload_seq = "text_encoder->text_encoder_2->image_encoder->unet->vae"
-    _optional_components = [
-        "tokenizer",
-        "tokenizer_2",
-        "text_encoder",
-        "text_encoder_2",
-        "image_encoder",
-        "feature_extractor",
-    ]
-    _callback_tensor_inputs = [
-        "latents",
-        "prompt_embeds",
-        "negative_prompt_embeds",
-        "add_text_embeds",
-        "add_time_ids",
-        "negative_pooled_prompt_embeds",
-        "negative_add_time_ids",
-    ]
-
-    def __init__(
-        self,
-        vae: AutoencoderKL,
-        text_encoder: CLIPTextModel,
-        text_encoder_2: CLIPTextModelWithProjection,
-        tokenizer: CLIPTokenizer,
-        tokenizer_2: CLIPTokenizer,
-        unet: UNet2DConditionModel,
-        scheduler: KarrasDiffusionSchedulers,
-        image_encoder: CLIPVisionModelWithProjection = None,
-        feature_extractor: CLIPImageProcessor = None,
-        force_zeros_for_empty_prompt: bool = True,
-        add_watermarker: Optional[bool] = None,
-    ):
-        super().__init__()
-
-        self.register_modules(
-            vae=vae,
-            text_encoder=text_encoder,
-            text_encoder_2=text_encoder_2,
-            tokenizer=tokenizer,
-            tokenizer_2=tokenizer_2,
-            unet=unet,
-            scheduler=scheduler,
-            image_encoder=image_encoder,
-            feature_extractor=feature_extractor,
-        )
-        self.register_to_config(force_zeros_for_empty_prompt=force_zeros_for_empty_prompt)
-        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
-        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
-
-        self.default_sample_size = self.unet.config.sample_size
-
-        add_watermarker = add_watermarker if add_watermarker is not None else is_invisible_watermark_available()
-
-        if add_watermarker:
-            self.watermark = StableDiffusionXLWatermarker()
-        else:
-            self.watermark = None
-
-    def encode_prompt(
-        self,
-        prompt: str,
-        prompt_2: Optional[str] = None,
-        device: Optional[torch.device] = None,
-        num_images_per_prompt: int = 1,
-        do_classifier_free_guidance: bool = True,
-        negative_prompt: Optional[str] = None,
-        negative_prompt_2: Optional[str] = None,
-        prompt_embeds: Optional[torch.Tensor] = None,
-        negative_prompt_embeds: Optional[torch.Tensor] = None,
-        pooled_prompt_embeds: Optional[torch.Tensor] = None,
-        negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
-        lora_scale: Optional[float] = None,
-        clip_skip: Optional[int] = None,
-    ):
-        device = device or self._execution_device
-
-        # set lora scale so that monkey patched LoRA
-        # function of text encoder can correctly access it
-        if lora_scale is not None and isinstance(self, StableDiffusionXLLoraLoaderMixin):
-            self._lora_scale = lora_scale
-
-            # dynamically adjust the LoRA scale
-            if self.text_encoder is not None:
-                if not USE_PEFT_BACKEND:
-                    adjust_lora_scale_text_encoder(self.text_encoder, lora_scale)
-                else:
-                    scale_lora_layers(self.text_encoder, lora_scale)
-
-            if self.text_encoder_2 is not None:
-                if not USE_PEFT_BACKEND:
-                    adjust_lora_scale_text_encoder(self.text_encoder_2, lora_scale)
-                else:
-                    scale_lora_layers(self.text_encoder_2, lora_scale)
-
-        prompt = [prompt] if isinstance(prompt, str) else prompt
-
-        if prompt is not None:
-            batch_size = len(prompt)
-        else:
-            batch_size = prompt_embeds.shape[0]
-
-        # Define tokenizers and text encoders
-        tokenizers = [self.tokenizer, self.tokenizer_2] if self.tokenizer is not None else [self.tokenizer_2]
-        text_encoders = (
-            [self.text_encoder, self.text_encoder_2] if self.text_encoder is not None else [self.text_encoder_2]
-        )
-
-        if prompt_embeds is None:
-            prompt_2 = prompt_2 or prompt
-            prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2
-
-            # textual inversion: process multi-vector tokens if necessary
-            prompt_embeds_list = []
-            prompts = [prompt, prompt_2]
-            for prompt, tokenizer, text_encoder in zip(prompts, tokenizers, text_encoders):
-                if isinstance(self, TextualInversionLoaderMixin):
-                    prompt = self.maybe_convert_prompt(prompt, tokenizer)
-
-                text_inputs = tokenizer(
-                    prompt,
-                    padding="max_length",
-                    max_length=tokenizer.model_max_length,
-                    truncation=True,
-                    return_tensors="pt",
-                )
-
-                text_input_ids = text_inputs.input_ids
-                untruncated_ids = tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
-
-                if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
-                    text_input_ids, untruncated_ids
-                ):
-                    removed_text = tokenizer.batch_decode(untruncated_ids[:, tokenizer.model_max_length - 1 : -1])
-                    logger.warning(
-                        "The following part of your input was truncated because CLIP can only handle sequences up to"
-                        f" {tokenizer.model_max_length} tokens: {removed_text}"
-                    )
-
-                prompt_embeds = text_encoder(text_input_ids.to(device), output_hidden_states=True)
-
-                # We are only ALWAYS interested in the pooled output of the final text encoder
-                pooled_prompt_embeds = prompt_embeds[0]
-                if clip_skip is None:
-                    prompt_embeds = prompt_embeds.hidden_states[-2]
-                else:
-                    # "2" because SDXL always indexes from the penultimate layer.
-                    prompt_embeds = prompt_embeds.hidden_states[-(clip_skip + 2)]
-
-                prompt_embeds_list.append(prompt_embeds)
-
-            prompt_embeds = torch.concat(prompt_embeds_list, dim=-1)
-
-        # get unconditional embeddings for classifier free guidance
-        zero_out_negative_prompt = negative_prompt is None and self.config.force_zeros_for_empty_prompt
-        if do_classifier_free_guidance and negative_prompt_embeds is None and zero_out_negative_prompt:
-            negative_prompt_embeds = torch.zeros_like(prompt_embeds)
-            negative_pooled_prompt_embeds = torch.zeros_like(pooled_prompt_embeds)
-        elif do_classifier_free_guidance and negative_prompt_embeds is None:
-            negative_prompt = negative_prompt or ""
-            negative_prompt_2 = negative_prompt_2 or negative_prompt
-
-            # normalize str to list
-            negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
-            negative_prompt_2 = (
-                batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2
-            )
-
-            uncond_tokens: List[str]
-            if prompt is not None and type(prompt) is not type(negative_prompt):
-                raise TypeError(
-                    f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
-                    f" {type(prompt)}."
-                )
-            elif batch_size != len(negative_prompt):
-                raise ValueError(
-                    f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
-                    f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
-                    " the batch size of `prompt`."
-                )
-            else:
-                uncond_tokens = [negative_prompt, negative_prompt_2]
-
-            negative_prompt_embeds_list = []
-            for negative_prompt, tokenizer, text_encoder in zip(uncond_tokens, tokenizers, text_encoders):
-                if isinstance(self, TextualInversionLoaderMixin):
-                    negative_prompt = self.maybe_convert_prompt(negative_prompt, tokenizer)
-
-                max_length = prompt_embeds.shape[1]
-                uncond_input = tokenizer(
-                    negative_prompt,
-                    padding="max_length",
-                    max_length=max_length,
-                    truncation=True,
-                    return_tensors="pt",
-                )
-
-                negative_prompt_embeds = text_encoder(
-                    uncond_input.input_ids.to(device),
-                    output_hidden_states=True,
-                )
-                # We are only ALWAYS interested in the pooled output of the final text encoder
-                negative_pooled_prompt_embeds = negative_prompt_embeds[0]
-                negative_prompt_embeds = negative_prompt_embeds.hidden_states[-2]
-
-                negative_prompt_embeds_list.append(negative_prompt_embeds)
-
-            negative_prompt_embeds = torch.concat(negative_prompt_embeds_list, dim=-1)
-
-        if self.text_encoder_2 is not None:
-            prompt_embeds = prompt_embeds.to(dtype=self.text_encoder_2.dtype, device=device)
-        else:
-            prompt_embeds = prompt_embeds.to(dtype=self.unet.dtype, device=device)
-
-        bs_embed, seq_len, _ = prompt_embeds.shape
-        # duplicate text embeddings for each generation per prompt, using mps friendly method
-        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
-        prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)
-
-        if do_classifier_free_guidance:
-            # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
-            seq_len = negative_prompt_embeds.shape[1]
-
-            if self.text_encoder_2 is not None:
-                negative_prompt_embeds = negative_prompt_embeds.to(dtype=self.text_encoder_2.dtype, device=device)
-            else:
-                negative_prompt_embeds = negative_prompt_embeds.to(dtype=self.unet.dtype, device=device)
-
-            negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
-            negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
-
-        pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt).view(
-            bs_embed * num_images_per_prompt, -1
-        )
-        if do_classifier_free_guidance:
-            negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.repeat(1, num_images_per_prompt).view(
-                bs_embed * num_images_per_prompt, -1
-            )
-
-        if self.text_encoder is not None:
-            if isinstance(self, StableDiffusionXLLoraLoaderMixin) and USE_PEFT_BACKEND:
-                # Retrieve the original scale by scaling back the LoRA layers
-                unscale_lora_layers(self.text_encoder, lora_scale)
-
-        if self.text_encoder_2 is not None:
-            if isinstance(self, StableDiffusionXLLoraLoaderMixin) and USE_PEFT_BACKEND:
-                # Retrieve the original scale by scaling back the LoRA layers
-                unscale_lora_layers(self.text_encoder_2, lora_scale)
-
-        return prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds
-
-    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.encode_image
-    def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None):
-        dtype = next(self.image_encoder.parameters()).dtype
-
-        if not isinstance(image, torch.Tensor):
-            image = self.feature_extractor(image, return_tensors="pt").pixel_values
-
-        image = image.to(device=device, dtype=dtype)
-        if output_hidden_states:
-            image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2]
-            image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0)
-            uncond_image_enc_hidden_states = self.image_encoder(
-                torch.zeros_like(image), output_hidden_states=True
-            ).hidden_states[-2]
-            uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave(
-                num_images_per_prompt, dim=0
-            )
-            return image_enc_hidden_states, uncond_image_enc_hidden_states
-        else:
-            image_embeds = self.image_encoder(image).image_embeds
-            image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0)
-            uncond_image_embeds = torch.zeros_like(image_embeds)
-
-            return image_embeds, uncond_image_embeds
-
-    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_ip_adapter_image_embeds
-    def prepare_ip_adapter_image_embeds(
-        self, ip_adapter_image, ip_adapter_image_embeds, device, num_images_per_prompt, do_classifier_free_guidance
-    ):
-        image_embeds = []
-        if do_classifier_free_guidance:
-            negative_image_embeds = []
-        if ip_adapter_image_embeds is None:
-            if not isinstance(ip_adapter_image, list):
-                ip_adapter_image = [ip_adapter_image]
-
-            if len(ip_adapter_image) != len(self.unet.encoder_hid_proj.image_projection_layers):
-                raise ValueError(
-                    f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(self.unet.encoder_hid_proj.image_projection_layers)} IP Adapters."
-                )
-
-            for single_ip_adapter_image, image_proj_layer in zip(
-                ip_adapter_image, self.unet.encoder_hid_proj.image_projection_layers
-            ):
-                output_hidden_state = not isinstance(image_proj_layer, ImageProjection)
-                single_image_embeds, single_negative_image_embeds = self.encode_image(
-                    single_ip_adapter_image, device, 1, output_hidden_state
-                )
-
-                image_embeds.append(single_image_embeds[None, :])
-                if do_classifier_free_guidance:
-                    negative_image_embeds.append(single_negative_image_embeds[None, :])
-        else:
-            for single_image_embeds in ip_adapter_image_embeds:
-                if do_classifier_free_guidance:
-                    single_negative_image_embeds, single_image_embeds = single_image_embeds.chunk(2)
-                    negative_image_embeds.append(single_negative_image_embeds)
-                image_embeds.append(single_image_embeds)
-
-        ip_adapter_image_embeds = []
-        for i, single_image_embeds in enumerate(image_embeds):
-            single_image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0)
-            if do_classifier_free_guidance:
-                single_negative_image_embeds = torch.cat([negative_image_embeds[i]] * num_images_per_prompt, dim=0)
-                single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds], dim=0)
-
-            single_image_embeds = single_image_embeds.to(device=device)
-            ip_adapter_image_embeds.append(single_image_embeds)
-
-        return ip_adapter_image_embeds
-
-    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs
-    def prepare_extra_step_kwargs(self, generator, eta):
-        # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
-        # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
-        # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
-        # and should be between [0, 1]
-
-        accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
-        extra_step_kwargs = {}
-        if accepts_eta:
-            extra_step_kwargs["eta"] = eta
-
-        # check if the scheduler accepts generator
-        accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
-        if accepts_generator:
-            extra_step_kwargs["generator"] = generator
-        return extra_step_kwargs
-
-    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents
-    def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None):
-        shape = (
-            batch_size,
-            num_channels_latents,
-            int(height) // self.vae_scale_factor,
-            int(width) // self.vae_scale_factor,
-        )
-        if isinstance(generator, list) and len(generator) != batch_size:
-            raise ValueError(
-                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
-                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
-            )
-
-        if latents is None:
-            latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
-        else:
-            latents = latents.to(device)
-
-        # scale the initial noise by the standard deviation required by the scheduler
-        latents = latents * self.scheduler.init_noise_sigma
-        return latents
-
-    def _get_add_time_ids(
-        self, original_size, crops_coords_top_left, target_size, dtype, text_encoder_projection_dim=None
-    ):
-        add_time_ids = list(original_size + crops_coords_top_left + target_size)
-
-        passed_add_embed_dim = (
-            self.unet.config.addition_time_embed_dim * len(add_time_ids) + text_encoder_projection_dim
-        )
-        expected_add_embed_dim = self.unet.add_embedding.linear_1.in_features
-
-        if expected_add_embed_dim != passed_add_embed_dim:
-            raise ValueError(
-                f"Model expects an added time embedding vector of length {expected_add_embed_dim}, but a vector of {passed_add_embed_dim} was created. The model has an incorrect config. Please check `unet.config.time_embedding_type` and `text_encoder_2.config.projection_dim`."
-            )
-
-        add_time_ids = torch.tensor([add_time_ids], dtype=dtype)
-        return add_time_ids
-
-    def upcast_vae(self):
-        dtype = self.vae.dtype
-        self.vae.to(dtype=torch.float32)
-        use_torch_2_0_or_xformers = isinstance(
-            self.vae.decoder.mid_block.attentions[0].processor,
-            (
-                AttnProcessor2_0,
-                XFormersAttnProcessor,
-                FusedAttnProcessor2_0,
-            ),
-        )
-        # if xformers or torch_2_0 is used attention block does not need
-        # to be in float32 which can save lots of memory
-        if use_torch_2_0_or_xformers:
-            self.vae.post_quant_conv.to(dtype)
-            self.vae.decoder.conv_in.to(dtype)
-            self.vae.decoder.mid_block.to(dtype)
-
-    # Copied from diffusers.pipelines.latent_consistency_models.pipeline_latent_consistency_text2img.LatentConsistencyModelPipeline.get_guidance_scale_embedding
-    def get_guidance_scale_embedding(
-        self, w: torch.Tensor, embedding_dim: int = 512, dtype: torch.dtype = torch.float32
-    ) -> torch.Tensor:
-        """
-        See https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298
-
-        Args:
-            w (`torch.Tensor`):
-                Generate embedding vectors with a specified guidance scale to subsequently enrich timestep embeddings.
-            embedding_dim (`int`, *optional*, defaults to 512):
-                Dimension of the embeddings to generate.
-            dtype (`torch.dtype`, *optional*, defaults to `torch.float32`):
-                Data type of the generated embeddings.
-
-        Returns:
-            `torch.Tensor`: Embedding vectors with shape `(len(w), embedding_dim)`.
-        """
-        assert len(w.shape) == 1
-        w = w * 1000.0
-
-        half_dim = embedding_dim // 2
-        emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1)
-        emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb)
-        emb = w.to(dtype)[:, None] * emb[None, :]
-        emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)
-        if embedding_dim % 2 == 1:  # zero pad
-            emb = torch.nn.functional.pad(emb, (0, 1))
-        assert emb.shape == (w.shape[0], embedding_dim)
-        return emb
-
-    @property
-    def guidance_scale(self):
-        return self._guidance_scale
-
-    @property
-    def guidance_rescale(self):
-        return self._guidance_rescale
-
-    @property
-    def clip_skip(self):
-        return self._clip_skip
-
-    # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
-    # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
-    # corresponds to doing no classifier free guidance.
-    @property
-    def do_classifier_free_guidance(self):
-        return self._guidance_scale > 1 and self.unet.config.time_cond_proj_dim is None
-
-    @property
-    def cross_attention_kwargs(self):
-        return self._cross_attention_kwargs
-
-    @property
-    def denoising_end(self):
-        return self._denoising_end
-
-    @property
-    def num_timesteps(self):
-        return self._num_timesteps
-
-    @property
-    def interrupt(self):
-        return self._interrupt
-
-    @torch.no_grad()
-    def __call__(
-        self,
-        prompt: Union[str, List[str]] = None,
-        prompt_2: Optional[Union[str, List[str]]] = None,
-        height: Optional[int] = None,
-        width: Optional[int] = None,
-        num_inference_steps: int = 50,
-        timesteps: List[int] = None,
-        sigmas: List[float] = None,
-        denoising_end: Optional[float] = None,
-        guidance_scale: float = 5.0,
-        negative_prompt: Optional[Union[str, List[str]]] = None,
-        negative_prompt_2: Optional[Union[str, List[str]]] = None,
-        num_images_per_prompt: Optional[int] = 1,
-        eta: float = 0.0,
-        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
-        latents: Optional[torch.Tensor] = None,
-        prompt_embeds: Optional[torch.Tensor] = None,
-        negative_prompt_embeds: Optional[torch.Tensor] = None,
-        pooled_prompt_embeds: Optional[torch.Tensor] = None,
-        negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
-        ip_adapter_image: Optional[PipelineImageInput] = None,
-        ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
-        output_type: Optional[str] = "pil",
-        return_dict: bool = True,
-        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
-        guidance_rescale: float = 0.0,
-        end_cfg: float = 0.9,
-        original_size: Optional[Tuple[int, int]] = None,
-        crops_coords_top_left: Tuple[int, int] = (0, 0),
-        target_size: Optional[Tuple[int, int]] = None,
-        negative_original_size: Optional[Tuple[int, int]] = None,
-        negative_crops_coords_top_left: Tuple[int, int] = (0, 0),
-        negative_target_size: Optional[Tuple[int, int]] = None,
-        clip_skip: Optional[int] = None,
-        callback_on_step_end: Optional[
-            Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
-        ] = None,
-        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
-        **kwargs,
-    ):
-        callback = kwargs.pop("callback", None)
-        callback_steps = kwargs.pop("callback_steps", None)
-
-        if callback is not None:
-            deprecate(
-                "callback",
-                "1.0.0",
-                "Passing `callback` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`",
-            )
-        if callback_steps is not None:
-            deprecate(
-                "callback_steps",
-                "1.0.0",
-                "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`",
-            )
-
-        if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)):
-            callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs
-
-        # 0. Default height and width to unet
-        height = height or self.default_sample_size * self.vae_scale_factor
-        width = width or self.default_sample_size * self.vae_scale_factor
-
-        original_size = original_size or (height, width)
-        target_size = target_size or (height, width)
-
-        self._guidance_scale = guidance_scale
-        self._guidance_rescale = guidance_rescale
-        self._clip_skip = clip_skip
-        self._cross_attention_kwargs = cross_attention_kwargs
-        self._denoising_end = denoising_end
-        self._interrupt = False
-
-        # 2. Define call parameters
-        if prompt is not None and isinstance(prompt, str):
-            batch_size = 1
-        elif prompt is not None and isinstance(prompt, list):
-            batch_size = len(prompt)
-        else:
-            batch_size = prompt_embeds.shape[0]
-
-        device = self._execution_device
-
-        # 3. Encode input prompt
-        lora_scale = (
-            self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None
-        )
-
-        (
-            prompt_embeds,
-            negative_prompt_embeds,
-            pooled_prompt_embeds,
-            negative_pooled_prompt_embeds,
-        ) = self.encode_prompt(
-            prompt=prompt,
-            prompt_2=prompt_2,
-            device=device,
-            num_images_per_prompt=num_images_per_prompt,
-            do_classifier_free_guidance=self.do_classifier_free_guidance,
-            negative_prompt=negative_prompt,
-            negative_prompt_2=negative_prompt_2,
-            prompt_embeds=prompt_embeds,
-            negative_prompt_embeds=negative_prompt_embeds,
-            pooled_prompt_embeds=pooled_prompt_embeds,
-            negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
-            lora_scale=lora_scale,
-            clip_skip=self.clip_skip,
-        )
-
-        # 4. Prepare timesteps
-        timesteps, num_inference_steps = retrieve_timesteps(
-            self.scheduler, num_inference_steps, device, timesteps, sigmas
-        )
-
-        # 5. Prepare latent variables
-        num_channels_latents = self.unet.config.in_channels
-        latents = self.prepare_latents(
-            batch_size * num_images_per_prompt,
-            num_channels_latents,
-            height,
-            width,
-            prompt_embeds.dtype,
-            device,
-            generator,
-            latents,
-        )
-
-        # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
-        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
-
-        # 7. Prepare added time ids & embeddings
-        add_text_embeds = pooled_prompt_embeds
-        if self.text_encoder_2 is None:
-            text_encoder_projection_dim = int(pooled_prompt_embeds.shape[-1])
-        else:
-            text_encoder_projection_dim = self.text_encoder_2.config.projection_dim
-
-        add_time_ids = self._get_add_time_ids(
-            original_size,
-            crops_coords_top_left,
-            target_size,
-            dtype=prompt_embeds.dtype,
-            text_encoder_projection_dim=text_encoder_projection_dim,
-        )
-        if negative_original_size is not None and negative_target_size is not None:
-            negative_add_time_ids = self._get_add_time_ids(
-                negative_original_size,
-                negative_crops_coords_top_left,
-                negative_target_size,
-                dtype=prompt_embeds.dtype,
-                text_encoder_projection_dim=text_encoder_projection_dim,
-            )
-        else:
-            negative_add_time_ids = add_time_ids
-
-        if self.do_classifier_free_guidance:
-            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
-            add_text_embeds = torch.cat([negative_pooled_prompt_embeds, add_text_embeds], dim=0)
-            add_time_ids = torch.cat([negative_add_time_ids, add_time_ids], dim=0)
-
-        prompt_embeds = prompt_embeds.to(device)
-        add_text_embeds = add_text_embeds.to(device)
-        add_time_ids = add_time_ids.to(device).repeat(batch_size * num_images_per_prompt, 1)
-
-        if ip_adapter_image is not None or ip_adapter_image_embeds is not None:
-            image_embeds = self.prepare_ip_adapter_image_embeds(
-                ip_adapter_image,
-                ip_adapter_image_embeds,
-                device,
-                batch_size * num_images_per_prompt,
-                self.do_classifier_free_guidance,
-            )
-
-        # 8. Denoising loop
-        num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
-
-        # 8.1 Apply denoising_end
-        if (
-            self.denoising_end is not None
-            and isinstance(self.denoising_end, float)
-            and self.denoising_end > 0
-            and self.denoising_end < 1
-        ):
-            discrete_timestep_cutoff = int(
-                round(
-                    self.scheduler.config.num_train_timesteps
-                    - (self.denoising_end * self.scheduler.config.num_train_timesteps)
-                )
-            )
-            num_inference_steps = len(list(filter(lambda ts: ts >= discrete_timestep_cutoff, timesteps)))
-            timesteps = timesteps[:num_inference_steps]
-
-        # 9. Optionally get Guidance Scale Embedding
-        timestep_cond = None
-        if self.unet.config.time_cond_proj_dim is not None:
-            guidance_scale_tensor = torch.tensor(self.guidance_scale - 1).repeat(batch_size * num_images_per_prompt)
-            timestep_cond = self.get_guidance_scale_embedding(
-                guidance_scale_tensor, embedding_dim=self.unet.config.time_cond_proj_dim
-            ).to(device=device, dtype=latents.dtype)
-
-        self._num_timesteps = len(timesteps)
-        with self.progress_bar(total=num_inference_steps) as progress_bar:
-            do_classifier_free_guidance = self.do_classifier_free_guidance
-            for i, t in enumerate(timesteps):
-                if self.interrupt:
-                    continue
-                if end_cfg is not None and i / num_inference_steps > end_cfg and do_classifier_free_guidance:
-                    do_classifier_free_guidance = False
-                    prompt_embeds = 1.5*torch.chunk(prompt_embeds, 2, dim=0)[-1]
-                    add_text_embeds = 1.5*torch.chunk(add_text_embeds, 2, dim=0)[-1]
-                    add_time_ids = 1.25*torch.chunk(add_time_ids, 2, dim=0)[-1]
-                # expand the latents if we are doing classifier free guidance
-                latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
-
-                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
-
-                # predict the noise residual
-                added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids}
-                if ip_adapter_image is not None or ip_adapter_image_embeds is not None:
-                    added_cond_kwargs["image_embeds"] = image_embeds
-                noise_pred = self.unet(
-                    latent_model_input,
-                    t,
-                    encoder_hidden_states=prompt_embeds,
-                    timestep_cond=timestep_cond,
-                    cross_attention_kwargs=self.cross_attention_kwargs,
-                    added_cond_kwargs=added_cond_kwargs,
-                    return_dict=False,
-                )[0]
-
-                # perform guidance
-                if do_classifier_free_guidance:
-                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
-                    noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)
-
-                if do_classifier_free_guidance and self.guidance_rescale > 0.0:
-                    # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf
-                    noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=self.guidance_rescale)
-
-                # compute the previous noisy sample x_t -> x_t-1
-                latents_dtype = latents.dtype
-                latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
-                if latents.dtype != latents_dtype:
-                    if torch.backends.mps.is_available():
-                        # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
-                        latents = latents.to(latents_dtype)
-
-                if callback_on_step_end is not None:
-                    callback_kwargs = {}
-                    for k in callback_on_step_end_tensor_inputs:
-                        callback_kwargs[k] = locals()[k]
-                    callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
-
-                    latents = callback_outputs.pop("latents", latents)
-                    prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
-                    negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds)
-                    add_text_embeds = callback_outputs.pop("add_text_embeds", add_text_embeds)
-                    negative_pooled_prompt_embeds = callback_outputs.pop(
-                        "negative_pooled_prompt_embeds", negative_pooled_prompt_embeds
-                    )
-                    add_time_ids = callback_outputs.pop("add_time_ids", add_time_ids)
-                    negative_add_time_ids = callback_outputs.pop("negative_add_time_ids", negative_add_time_ids)
-
-                # call the callback, if provided
-                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
-                    progress_bar.update()
-                    if callback is not None and i % callback_steps == 0:
-                        step_idx = i // getattr(self.scheduler, "order", 1)
-                        callback(step_idx, t, latents)
-
-                if XLA_AVAILABLE:
-                    xm.mark_step()
-
-        if not output_type == "latent":
-            # make sure the VAE is in float32 mode, as it overflows in float16
-            needs_upcasting = self.vae.dtype == torch.float16 and self.vae.config.force_upcast
-
-            if needs_upcasting:
-                self.upcast_vae()
-                latents = latents.to(next(iter(self.vae.post_quant_conv.parameters())).dtype)
-            elif latents.dtype != self.vae.dtype:
-                if torch.backends.mps.is_available():
-                    # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
-                    self.vae = self.vae.to(latents.dtype)
-
-            # unscale/denormalize the latents
-            # denormalize with the mean and std if available and not None
-            has_latents_mean = hasattr(self.vae.config, "latents_mean") and self.vae.config.latents_mean is not None
-            has_latents_std = hasattr(self.vae.config, "latents_std") and self.vae.config.latents_std is not None
-            if has_latents_mean and has_latents_std:
-                latents_mean = (
-                    torch.tensor(self.vae.config.latents_mean).view(1, 4, 1, 1).to(latents.device, latents.dtype)
-                )
-                latents_std = (
-                    torch.tensor(self.vae.config.latents_std).view(1, 4, 1, 1).to(latents.device, latents.dtype)
-                )
-                latents = latents * latents_std / self.vae.config.scaling_factor + latents_mean
-            else:
-                latents = latents / self.vae.config.scaling_factor
-
-            image = self.vae.decode(latents, return_dict=False)[0]
-
-            # cast back to fp16 if needed
-            if needs_upcasting:
-                self.vae.to(dtype=torch.float16)
-        else:
-            image = latents
-
-        if not output_type == "latent":
-            # apply watermark if available
-            if self.watermark is not None:
-                image = self.watermark.apply_watermark(image)
-
-            image = self.image_processor.postprocess(image, output_type=output_type)
-
-        # Offload all models
-        self.maybe_free_model_hooks()
-
-        if not return_dict:
-            return (image,)
-
-        return StableDiffusionXLPipelineOutput(images=image)
-
-def load_pipeline(pipeline=None) -> StableDiffusionXLPipeline:
-    """Load and prepare the pipeline."""
-    if not pipeline:
-        pipeline = StableDiffusionXLPipeline_new.from_pretrained(
-            "./models/newdream-sdxl-20",
-            torch_dtype=torch.float16,
-            local_files_only=True,
-        ).to("cuda")
-    
-    pipeline.scheduler = SchedulerWrapper(DDIMScheduler.from_config(pipeline.scheduler.config))
-    pipeline = compile_pipe(pipeline)
-    load_pipe(pipeline, dir="cached_pipe")
-    
-    # Warm-up runs
-    for _ in range(4):
-        pipeline(
-            prompt="a cat and a dog",
-            num_inference_steps=18
-        )
-    pipeline.scheduler.prepare_loss()
-    return pipeline
 
+generator = Generator(torch.device("cuda")).manual_seed(69)
+prompt = "make submissions great again"
+SDXL_DEFAULT_CONFIG = [
+            {
+                "wildcard_or_filter_func": lambda name: "down_blocks.2" not in name and"down_blocks.3" not in name and "up_blocks.2" not in name,
+                "select_cache_step_func": lambda step: (step % 2 != 0) and (step >= 10),
+                }]
+def load_pipeline() -> StableDiffusionXLPipeline:
+    pipe = StableDiffusionXLPipeline.from_pretrained(
+        "models/newdream-sdxl-20", torch_dtype=torch.float16, use_safetensors=True, local_files_only=True
+    ).to("cuda") 
+    load_unet_trt(
+        pipe.unet,
+        engine_path=Path("./engine"),
+        batch_size=1,
+    )
+    cachify.prepare(pipe, SDXL_DEFAULT_CONFIG)
+    cachify.enable(pipe)
+    pipe.scheduler = SchedulerWrapper(DDIMScheduler.from_config(pipe.scheduler.config))
+    with cachify.infer(pipe) as cached_pipe:
+        for _ in range(5):
+            pipe(prompt=prompt, num_inference_steps=20)
+    cachify.disable(pipe)
+    pipe.scheduler.prepare_loss()
+    return pipe
 
 def infer(request: TextToImageRequest, pipeline: StableDiffusionXLPipeline) -> Image:
-    """Generate image from text prompt."""
-    generator = Generator(pipeline.device).manual_seed(request.seed) if request.seed else None
-    
-    image_0 = pipeline(
-        prompt=request.prompt,
-        negative_prompt=request.negative_prompt,
-        width=request.width,
-        height=request.height,
-        generator=generator,
-        num_inference_steps=18,
-    ).images[0]
-
-    filter_image = max_pixel_filter(image_0)
-    return filter_image
-
 
+    if request.seed is None:
+        generator = None
+    else:
+        generator = Generator(pipeline.device).manual_seed(request.seed)
+    cachify.prepare(pipeline, SDXL_DEFAULT_CONFIG)
+    cachify.enable(pipeline)
+    with cachify.infer(pipeline) as cached_pipe:
+        image = cached_pipe(
+            prompt=request.prompt,
+            negative_prompt=request.negative_prompt,
+            width=request.width,
+            height=request.height,
+            generator=generator,
+            num_inference_steps=15,
+            ).images[0]
+    return image

src/trt_pipeline/config.py

@@ -0,0 +1,162 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: MIT
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+from diffusers.models.transformers.transformer_sd3 import SD3Transformer2DModel
+from diffusers.models.unets.unet_2d_condition import UNet2DConditionModel
+
+sd3_common_transformer_block_config = {
+    "dummy_input": {
+        "hidden_states": (2, 4096, 1536),
+        "encoder_hidden_states": (2, 333, 1536),
+        "temb": (2, 1536),
+    },
+    "output_names": ["encoder_hidden_states_out", "hidden_states_out"],
+    "dynamic_axes": {
+        "hidden_states": {0: "batch_size"},
+        "encoder_hidden_states": {0: "batch_size"},
+        "temb": {0: "steps"},
+    },
+}
+
+ONNX_CONFIG = {
+    UNet2DConditionModel: {
+        "down_blocks.0": {
+            "dummy_input": {
+                "hidden_states": (2, 320, 128, 128),
+                "temb": (2, 1280),
+            },
+            "output_names": ["sample", "res_samples_0", "res_samples_1", "res_samples_2"],
+            "dynamic_axes": {
+                "hidden_states": {0: "batch_size"},
+                "temb": {0: "steps"},
+            },
+        },
+        "down_blocks.1": {
+            "dummy_input": {
+                "hidden_states": (2, 320, 64, 64),
+                "temb": (2, 1280),
+                "encoder_hidden_states": (2, 77, 2048),
+            },
+            "output_names": ["sample", "res_samples_0", "res_samples_1", "res_samples_2"],
+            "dynamic_axes": {
+                "hidden_states": {0: "batch_size"},
+                "temb": {0: "steps"},
+                "encoder_hidden_states": {0: "batch_size"},
+            },
+        },
+        "down_blocks.2": {
+            "dummy_input": {
+                "hidden_states": (2, 640, 32, 32),
+                "temb": (2, 1280),
+                "encoder_hidden_states": (2, 77, 2048),
+            },
+            "output_names": ["sample", "res_samples_0", "res_samples_1"],
+            "dynamic_axes": {
+                "hidden_states": {0: "batch_size"},
+                "temb": {0: "steps"},
+                "encoder_hidden_states": {0: "batch_size"},
+            },
+        },
+        "mid_block": {
+            "dummy_input": {
+                "hidden_states": (2, 1280, 32, 32),
+                "temb": (2, 1280),
+                "encoder_hidden_states": (2, 77, 2048),
+            },
+            "output_names": ["sample"],
+            "dynamic_axes": {
+                "hidden_states": {0: "batch_size"},
+                "temb": {0: "steps"},
+                "encoder_hidden_states": {0: "batch_size"},
+            },
+        },
+        "up_blocks.0": {
+            "dummy_input": {
+                "hidden_states": (2, 1280, 32, 32),
+                "res_hidden_states_0": (2, 640, 32, 32),
+                "res_hidden_states_1": (2, 1280, 32, 32),
+                "res_hidden_states_2": (2, 1280, 32, 32),
+                "temb": (2, 1280),
+                "encoder_hidden_states": (2, 77, 2048),
+            },
+            "output_names": ["sample"],
+            "dynamic_axes": {
+                "hidden_states": {0: "batch_size"},
+                "temb": {0: "steps"},
+                "encoder_hidden_states": {0: "batch_size"},
+                "res_hidden_states_0": {0: "batch_size"},
+                "res_hidden_states_1": {0: "batch_size"},
+                "res_hidden_states_2": {0: "batch_size"},
+            },
+        },
+        "up_blocks.1": {
+            "dummy_input": {
+                "hidden_states": (2, 1280, 64, 64),
+                "res_hidden_states_0": (2, 320, 64, 64),
+                "res_hidden_states_1": (2, 640, 64, 64),
+                "res_hidden_states_2": (2, 640, 64, 64),
+                "temb": (2, 1280),
+                "encoder_hidden_states": (2, 77, 2048),
+            },
+            "output_names": ["sample"],
+            "dynamic_axes": {
+                "hidden_states": {0: "batch_size"},
+                "temb": {0: "steps"},
+                "encoder_hidden_states": {0: "batch_size"},
+                "res_hidden_states_0": {0: "batch_size"},
+                "res_hidden_states_1": {0: "batch_size"},
+                "res_hidden_states_2": {0: "batch_size"},
+            },
+        },
+        "up_blocks.2": {
+            "dummy_input": {
+                "hidden_states": (2, 640, 128, 128),
+                "res_hidden_states_0": (2, 320, 128, 128),
+                "res_hidden_states_1": (2, 320, 128, 128),
+                "res_hidden_states_2": (2, 320, 128, 128),
+                "temb": (2, 1280),
+            },
+            "output_names": ["sample"],
+            "dynamic_axes": {
+                "hidden_states": {0: "batch_size"},
+                "temb": {0: "steps"},
+                "res_hidden_states_0": {0: "batch_size"},
+                "res_hidden_states_1": {0: "batch_size"},
+                "res_hidden_states_2": {0: "batch_size"},
+            },
+        },
+    },
+    SD3Transformer2DModel: {
+        **{f"transformer_blocks.{i}": sd3_common_transformer_block_config for i in range(23)},
+        "transformer_blocks.23": {
+            "dummy_input": {
+                "hidden_states": (2, 4096, 1536),
+                "encoder_hidden_states": (2, 333, 1536),
+                "temb": (2, 1536),
+            },
+            "output_names": ["hidden_states_out"],
+            "dynamic_axes": {
+                "hidden_states": {0: "batch_size"},
+                "encoder_hidden_states": {0: "batch_size"},
+                "temb": {0: "steps"},
+            },
+        },
+    },
+}

src/trt_pipeline/deploy.py

@@ -0,0 +1,144 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: MIT
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+
+import types
+from pathlib import Path
+
+import tensorrt as trt
+import torch
+from cache_diffusion.cachify import CACHED_PIPE, get_model
+from cuda import cudart
+from diffusers.models.transformers.transformer_sd3 import SD3Transformer2DModel
+from diffusers.models.unets.unet_2d_condition import UNet2DConditionModel
+from trt_pipeline.config import ONNX_CONFIG
+from trt_pipeline.models.sd3 import sd3_forward
+from trt_pipeline.models.sdxl import (
+    cachecrossattnupblock2d_forward,
+    cacheunet_forward,
+    cacheupblock2d_forward,
+)
+from polygraphy.backend.trt import (
+    CreateConfig,
+    Profile,
+    engine_from_network,
+    network_from_onnx_path,
+    save_engine,
+)
+from torch.onnx import export as onnx_export
+
+from .utils import Engine
+
+
+def replace_new_forward(backbone):
+    if backbone.__class__ == UNet2DConditionModel:
+        backbone.forward = types.MethodType(cacheunet_forward, backbone)
+        for upsample_block in backbone.up_blocks:
+            if (
+                hasattr(upsample_block, "has_cross_attention")
+                and upsample_block.has_cross_attention
+            ):
+                upsample_block.forward = types.MethodType(
+                    cachecrossattnupblock2d_forward, upsample_block
+                )
+            else:
+                upsample_block.forward = types.MethodType(cacheupblock2d_forward, upsample_block)
+    elif backbone.__class__ == SD3Transformer2DModel:
+        backbone.forward = types.MethodType(sd3_forward, backbone)
+
+
+def get_input_info(dummy_dict, info: str = None, batch_size: int = 1):
+    return_val = [] if info == "profile_shapes" or info == "input_names" else {}
+
+    def collect_leaf_keys(d):
+        for key, value in d.items():
+            if isinstance(value, dict):
+                collect_leaf_keys(value)
+            else:
+                value = (value[0] * batch_size,) + value[1:]
+                if info == "profile_shapes":
+                    return_val.append((key, value))  # type: ignore
+                elif info == "profile_shapes_dict":
+                    return_val[key] = value  # type: ignore
+                elif info == "dummy_input":
+                    return_val[key] = torch.ones(value).half().cuda()  # type: ignore
+                elif info == "input_names":
+                    return_val.append(key)  # type: ignore
+
+    collect_leaf_keys(dummy_dict)
+    return return_val
+
+
+def get_total_device_memory(backbone):
+    max_device_memory = 0
+    for _, engine in backbone.engines.items():
+        max_device_memory = max(max_device_memory, engine.engine.device_memory_size)
+    return max_device_memory
+
+
+def load_engines(backbone, engine_path: Path, batch_size: int = 1):
+    backbone.engines = {}
+    for f in engine_path.iterdir():
+        if f.is_file():
+            eng = Engine()
+            eng.load(str(f))
+            backbone.engines[f"{f.stem}"] = eng
+    _, shared_device_memory = cudart.cudaMalloc(get_total_device_memory(backbone))
+    for engine in backbone.engines.values():
+        engine.activate(shared_device_memory)
+    backbone.cuda_stream = cudart.cudaStreamCreate()[1]
+    for block_name in backbone.engines.keys():
+        backbone.engines[block_name].allocate_buffers(
+            shape_dict=get_input_info(
+                ONNX_CONFIG[backbone.__class__][block_name]["dummy_input"],
+                "profile_shapes_dict",
+                batch_size,
+            ),
+            device=backbone.device,
+            batch_size=batch_size,
+        )
+    # TODO: Free and clean up the origin pytorch cuda memory
+
+
+def warm_up(backbone, batch_size: int = 1):
+    print("Warming-up TensorRT engines...")
+    for name, engine in backbone.engines.items():
+        dummy_input = get_input_info(
+            ONNX_CONFIG[backbone.__class__][name]["dummy_input"], "dummy_input", batch_size
+        )
+        _ = engine(dummy_input, backbone.cuda_stream)
+
+
+def teardown(pipe):
+    backbone = get_model(pipe)
+    for engine in backbone.engines.values():
+        del engine
+
+    cudart.cudaStreamDestroy(backbone.cuda_stream)
+    del backbone.cuda_stream
+
+
+def load_unet_trt(unet, engine_path: Path, batch_size: int = 1):
+    backbone = unet 
+    engine_path.mkdir(parents=True, exist_ok=True)
+    replace_new_forward(backbone) 
+    load_engines(backbone, engine_path, batch_size)
+    warm_up(backbone, batch_size)
+    backbone.use_trt_infer = True

src/trt_pipeline/models/sd3.py

@@ -0,0 +1,159 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: MIT
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+
+from typing import Any, Dict, List, Optional, Union
+
+import torch
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.utils import (
+    USE_PEFT_BACKEND,
+    is_torch_version,
+    scale_lora_layers,
+    unscale_lora_layers,
+)
+
+
+def sd3_forward(
+    self,
+    hidden_states: torch.FloatTensor,
+    encoder_hidden_states: torch.FloatTensor = None,
+    pooled_projections: torch.FloatTensor = None,
+    timestep: torch.LongTensor = None,
+    block_controlnet_hidden_states: List = None,
+    joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+    return_dict: bool = True,
+) -> Union[torch.FloatTensor, Transformer2DModelOutput]:
+    """
+    The [`SD3Transformer2DModel`] forward method.
+
+    Args:
+        hidden_states (`torch.FloatTensor` of shape `(batch size, channel, height, width)`):
+            Input `hidden_states`.
+        encoder_hidden_states (`torch.FloatTensor` of shape `(batch size, sequence_len, embed_dims)`):
+            Conditional embeddings (embeddings computed from the input conditions such as prompts) to use.
+        pooled_projections (`torch.FloatTensor` of shape `(batch_size, projection_dim)`): Embeddings projected
+            from the embeddings of input conditions.
+        timestep ( `torch.LongTensor`):
+            Used to indicate denoising step.
+        block_controlnet_hidden_states: (`list` of `torch.Tensor`):
+            A list of tensors that if specified are added to the residuals of transformer blocks.
+        joint_attention_kwargs (`dict`, *optional*):
+            A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+            `self.processor` in
+            [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+        return_dict (`bool`, *optional*, defaults to `True`):
+            Whether or not to return a [`~models.transformer_2d.Transformer2DModelOutput`] instead of a plain
+            tuple.
+
+    Returns:
+        If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a
+        `tuple` where the first element is the sample tensor.
+    """
+    if joint_attention_kwargs is not None:
+        joint_attention_kwargs = joint_attention_kwargs.copy()
+        lora_scale = joint_attention_kwargs.pop("scale", 1.0)
+    else:
+        lora_scale = 1.0
+
+    if USE_PEFT_BACKEND:
+        # weight the lora layers by setting `lora_scale` for each PEFT layer
+        scale_lora_layers(self, lora_scale)
+
+    height, width = hidden_states.shape[-2:]
+
+    hidden_states = self.pos_embed(hidden_states)  # takes care of adding positional embeddings too.
+    temb = self.time_text_embed(timestep, pooled_projections)
+    encoder_hidden_states = self.context_embedder(encoder_hidden_states)
+
+    for index_block, block in enumerate(self.transformer_blocks):
+        if self.training and self.gradient_checkpointing:
+
+            def create_custom_forward(module, return_dict=None):
+                def custom_forward(*inputs):
+                    if return_dict is not None:
+                        return module(*inputs, return_dict=return_dict)
+                    else:
+                        return module(*inputs)
+
+                return custom_forward
+
+            ckpt_kwargs: Dict[str, Any] = (
+                {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+            )
+            encoder_hidden_states, hidden_states = torch.utils.checkpoint.checkpoint(
+                create_custom_forward(block),
+                hidden_states,
+                encoder_hidden_states,
+                temb,
+                **ckpt_kwargs,
+            )
+
+        else:
+            if hasattr(self, "use_trt_infer") and self.use_trt_infer:
+                feed_dict = {
+                    "hidden_states": hidden_states,
+                    "encoder_hidden_states": encoder_hidden_states,
+                    "temb": temb,
+                }
+                _results = self.engines[f"transformer_blocks.{index_block}"](
+                    feed_dict, self.cuda_stream
+                )
+                if index_block != 23:
+                    encoder_hidden_states = _results["encoder_hidden_states_out"]
+                hidden_states = _results["hidden_states_out"]
+            else:
+                encoder_hidden_states, hidden_states = block(
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    temb=temb,
+                )
+
+        # controlnet residual
+        if block_controlnet_hidden_states is not None and block.context_pre_only is False:
+            interval_control = len(self.transformer_blocks) // len(block_controlnet_hidden_states)
+            hidden_states = (
+                hidden_states + block_controlnet_hidden_states[index_block // interval_control]
+            )
+
+    hidden_states = self.norm_out(hidden_states, temb)
+    hidden_states = self.proj_out(hidden_states)
+
+    # unpatchify
+    patch_size = self.config.patch_size
+    height = height // patch_size
+    width = width // patch_size
+
+    hidden_states = hidden_states.reshape(
+        shape=(hidden_states.shape[0], height, width, patch_size, patch_size, self.out_channels)
+    )
+    hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states)
+    output = hidden_states.reshape(
+        shape=(hidden_states.shape[0], self.out_channels, height * patch_size, width * patch_size)
+    )
+
+    if USE_PEFT_BACKEND:
+        # remove `lora_scale` from each PEFT layer
+        unscale_lora_layers(self, lora_scale)
+
+    if not return_dict:
+        return (output,)
+
+    return Transformer2DModelOutput(sample=output)

src/trt_pipeline/models/sdxl.py

@@ -0,0 +1,275 @@
+# Adapted from
+# https://github.com/huggingface/diffusers/blob/73acebb8cfbd1d2954cabe1af4185f9994e61917/src/diffusers/models/unets/unet_2d_condition.py#L1039-L1312
+# https://github.com/huggingface/diffusers/blob/73acebb8cfbd1d2954cabe1af4185f9994e61917/src/diffusers/models/unets/unet_2d_blocks.py#L2482-L2564
+# https://github.com/huggingface/diffusers/blob/73acebb8cfbd1d2954cabe1af4185f9994e61917/src/diffusers/models/unets/unet_2d_blocks.py#L2617-L2679
+
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# Not a contribution
+# Changes made by NVIDIA CORPORATION & AFFILIATES or otherwise documented as
+# NVIDIA-proprietary are not a contribution and subject to the following terms and conditions:
+# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
+#
+# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
+# property and proprietary rights in and to this material, related
+# documentation and any modifications thereto. Any use, reproduction,
+# disclosure or distribution of this material and related documentation
+# without an express license agreement from NVIDIA CORPORATION or
+# its affiliates is strictly prohibited.
+
+from typing import Any, Dict, Optional, Tuple, Union
+
+import torch
+from diffusers.models.unets.unet_2d_condition import UNet2DConditionOutput
+
+
+def cachecrossattnupblock2d_forward(
+    self,
+    hidden_states: torch.FloatTensor,
+    res_hidden_states_0: torch.FloatTensor,
+    res_hidden_states_1: torch.FloatTensor,
+    res_hidden_states_2: torch.FloatTensor,
+    temb: Optional[torch.FloatTensor] = None,
+    encoder_hidden_states: Optional[torch.FloatTensor] = None,
+    cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+    upsample_size: Optional[int] = None,
+    attention_mask: Optional[torch.FloatTensor] = None,
+    encoder_attention_mask: Optional[torch.FloatTensor] = None,
+) -> torch.FloatTensor:
+    res_hidden_states_tuple = (res_hidden_states_0, res_hidden_states_1, res_hidden_states_2)
+    for resnet, attn in zip(self.resnets, self.attentions):
+        # pop res hidden states
+        res_hidden_states = res_hidden_states_tuple[-1]
+        res_hidden_states_tuple = res_hidden_states_tuple[:-1]
+
+        hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
+
+        hidden_states = resnet(hidden_states, temb)
+        hidden_states = attn(
+            hidden_states,
+            encoder_hidden_states=encoder_hidden_states,
+            cross_attention_kwargs=cross_attention_kwargs,
+            attention_mask=attention_mask,
+            encoder_attention_mask=encoder_attention_mask,
+            return_dict=False,
+        )[0]
+
+    if self.upsamplers is not None:
+        for upsampler in self.upsamplers:
+            hidden_states = upsampler(hidden_states, upsample_size)
+
+    return hidden_states
+
+
+def cacheupblock2d_forward(
+    self,
+    hidden_states: torch.FloatTensor,
+    res_hidden_states_0: torch.FloatTensor,
+    res_hidden_states_1: torch.FloatTensor,
+    res_hidden_states_2: torch.FloatTensor,
+    temb: Optional[torch.FloatTensor] = None,
+    upsample_size: Optional[int] = None,
+) -> torch.FloatTensor:
+    res_hidden_states_tuple = (res_hidden_states_0, res_hidden_states_1, res_hidden_states_2)
+    for resnet in self.resnets:
+        # pop res hidden states
+        res_hidden_states = res_hidden_states_tuple[-1]
+        res_hidden_states_tuple = res_hidden_states_tuple[:-1]
+
+        hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
+
+        hidden_states = resnet(hidden_states, temb)
+
+    if self.upsamplers is not None:
+        for upsampler in self.upsamplers:
+            hidden_states = upsampler(hidden_states, upsample_size)
+
+    return hidden_states
+
+
+def cacheunet_forward(
+    self,
+    sample: torch.FloatTensor,
+    timestep: Union[torch.Tensor, float, int],
+    encoder_hidden_states: torch.Tensor,
+    class_labels: Optional[torch.Tensor] = None,
+    timestep_cond: Optional[torch.Tensor] = None,
+    attention_mask: Optional[torch.Tensor] = None,
+    cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+    added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None,
+    down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None,
+    mid_block_additional_residual: Optional[torch.Tensor] = None,
+    down_intrablock_additional_residuals: Optional[Tuple[torch.Tensor]] = None,
+    encoder_attention_mask: Optional[torch.Tensor] = None,
+    return_dict: bool = True,
+) -> Union[UNet2DConditionOutput, Tuple]:
+    # 1. time
+    t_emb = self.get_time_embed(sample=sample, timestep=timestep)
+    emb = self.time_embedding(t_emb, timestep_cond)
+    aug_emb = None
+
+    aug_emb = self.get_aug_embed(
+        emb=emb,
+        encoder_hidden_states=encoder_hidden_states,
+        added_cond_kwargs=added_cond_kwargs,
+    )
+
+    emb = emb + aug_emb if aug_emb is not None else emb
+
+    encoder_hidden_states = self.process_encoder_hidden_states(
+        encoder_hidden_states=encoder_hidden_states, added_cond_kwargs=added_cond_kwargs
+    )
+
+    # 2. pre-process
+    sample = self.conv_in(sample)
+
+    if hasattr(self, "_export_precess_onnx") and self._export_precess_onnx:
+        return (
+            sample,
+            encoder_hidden_states,
+            emb,
+        )
+
+    down_block_res_samples = (sample,)
+    for i, downsample_block in enumerate(self.down_blocks):
+        if (
+            hasattr(downsample_block, "has_cross_attention")
+            and downsample_block.has_cross_attention
+        ):
+            if hasattr(self, "use_trt_infer") and self.use_trt_infer:
+                feed_dict = {
+                    "hidden_states": sample,
+                    "temb": emb,
+                    "encoder_hidden_states": encoder_hidden_states,
+                }
+                down_results = self.engines[f"down_blocks.{i}"](feed_dict, self.cuda_stream)
+                sample = down_results["sample"]
+                res_samples_0 = down_results["res_samples_0"]
+                res_samples_1 = down_results["res_samples_1"]
+                if "res_samples_2" in down_results.keys():
+                    res_samples_2 = down_results["res_samples_2"]
+            else:
+                # For t2i-adapter CrossAttnDownBlock2D
+                additional_residuals = {}
+
+                sample, res_samples = downsample_block(
+                    hidden_states=sample,
+                    temb=emb,
+                    encoder_hidden_states=encoder_hidden_states,
+                    attention_mask=attention_mask,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                    encoder_attention_mask=encoder_attention_mask,
+                    **additional_residuals,
+                )
+        else:
+            if hasattr(self, "use_trt_infer") and self.use_trt_infer:
+                feed_dict = {"hidden_states": sample, "temb": emb}
+                down_results = self.engines[f"down_blocks.{i}"](feed_dict, self.cuda_stream)
+                sample = down_results["sample"]
+                res_samples_0 = down_results["res_samples_0"]
+                res_samples_1 = down_results["res_samples_1"]
+                if "res_samples_2" in down_results.keys():
+                    res_samples_2 = down_results["res_samples_2"]
+            else:
+                sample, res_samples = downsample_block(hidden_states=sample, temb=emb)
+
+        if hasattr(self, "use_trt_infer") and self.use_trt_infer:
+            down_block_res_samples += (
+                res_samples_0,
+                res_samples_1,
+            )
+            if "res_samples_2" in down_results.keys():
+                down_block_res_samples += (res_samples_2,)
+        else:
+            down_block_res_samples += res_samples
+
+    if hasattr(self, "use_trt_infer") and self.use_trt_infer:
+        feed_dict = {
+            "hidden_states": sample,
+            "temb": emb,
+            "encoder_hidden_states": encoder_hidden_states,
+        }
+        mid_results = self.engines["mid_block"](feed_dict, self.cuda_stream)
+        sample = mid_results["sample"]
+    else:
+        sample = self.mid_block(
+            sample,
+            emb,
+            encoder_hidden_states=encoder_hidden_states,
+            attention_mask=attention_mask,
+            cross_attention_kwargs=cross_attention_kwargs,
+            encoder_attention_mask=encoder_attention_mask,
+        )
+
+    # 5. up
+    for i, upsample_block in enumerate(self.up_blocks):
+        res_samples = down_block_res_samples[-len(upsample_block.resnets) :]
+        down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]
+
+        if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention:
+            if hasattr(self, "use_trt_infer") and self.use_trt_infer:
+                feed_dict = {
+                    "hidden_states": sample,
+                    "res_hidden_states_0": res_samples[0],
+                    "res_hidden_states_1": res_samples[1],
+                    "res_hidden_states_2": res_samples[2],
+                    "temb": emb,
+                    "encoder_hidden_states": encoder_hidden_states,
+                }
+                up_results = self.engines[f"up_blocks.{i}"](feed_dict, self.cuda_stream)
+                sample = up_results["sample"]
+            else:
+                sample = upsample_block(
+                    hidden_states=sample,
+                    temb=emb,
+                    res_hidden_states_0=res_samples[0],
+                    res_hidden_states_1=res_samples[1],
+                    res_hidden_states_2=res_samples[2],
+                    encoder_hidden_states=encoder_hidden_states,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                    attention_mask=attention_mask,
+                    encoder_attention_mask=encoder_attention_mask,
+                )
+        else:
+            if hasattr(self, "use_trt_infer") and self.use_trt_infer:
+                feed_dict = {
+                    "hidden_states": sample,
+                    "res_hidden_states_0": res_samples[0],
+                    "res_hidden_states_1": res_samples[1],
+                    "res_hidden_states_2": res_samples[2],
+                    "temb": emb,
+                }
+                up_results = self.engines[f"up_blocks.{i}"](feed_dict, self.cuda_stream)
+                sample = up_results["sample"]
+            else:
+                sample = upsample_block(
+                    hidden_states=sample,
+                    temb=emb,
+                    res_hidden_states_0=res_samples[0],
+                    res_hidden_states_1=res_samples[1],
+                    res_hidden_states_2=res_samples[2],
+                )
+
+    # 6. post-process
+    if self.conv_norm_out:
+        sample = self.conv_norm_out(sample)
+        sample = self.conv_act(sample)
+    sample = self.conv_out(sample)
+
+    if not return_dict:
+        return (sample,)
+
+    return UNet2DConditionOutput(sample=sample)

src/trt_pipeline/utils.py

@@ -0,0 +1,129 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: MIT
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+
+from collections import OrderedDict
+
+import numpy as np
+import tensorrt as trt
+import torch
+from cuda import cudart
+from polygraphy.backend.common import bytes_from_path
+from polygraphy.backend.trt import engine_from_bytes
+
+numpy_to_torch_dtype_dict = {
+    np.uint8: torch.uint8,
+    np.int8: torch.int8,
+    np.int16: torch.int16,
+    np.int32: torch.int32,
+    np.int64: torch.int64,
+    np.float16: torch.float16,
+    np.float32: torch.float32,
+    np.float64: torch.float64,
+    np.complex64: torch.complex64,
+    np.complex128: torch.complex128,
+}
+
+
+class Engine:
+    def __init__(
+        self,
+    ):
+        self.engine = None
+        self.context = None
+        self.buffers = OrderedDict()
+        self.tensors = OrderedDict()
+        self.cuda_graph_instance = None  # cuda graph
+        self.has_cross_attention = False
+
+    def __del__(self):
+        del self.engine
+        del self.context
+        del self.buffers
+        del self.tensors
+
+    def load(self, engine_path):
+        self.engine = engine_from_bytes(bytes_from_path(engine_path))
+
+    def activate(self, reuse_device_memory=None):
+        if reuse_device_memory:
+            self.context = self.engine.create_execution_context_without_device_memory()  # type: ignore
+            self.context.device_memory = reuse_device_memory
+        else:
+            self.context = self.engine.create_execution_context()  # type: ignore
+
+    def allocate_buffers(self, shape_dict=None, device="cuda", batch_size=1):
+        for binding in range(self.engine.num_io_tensors):  # type: ignore
+            name = self.engine.get_tensor_name(binding)  # type: ignore
+            if shape_dict and name in shape_dict:
+                shape = shape_dict[name]
+            else:
+                shape = self.engine.get_tensor_shape(name)  # type: ignore
+                shape = (batch_size * 2,) + shape[1:]
+            dtype = trt.nptype(self.engine.get_tensor_dtype(name))  # type: ignore
+            if self.engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT:  # type: ignore
+                self.context.set_input_shape(name, shape)  # type: ignore
+            tensor = torch.empty(tuple(shape), dtype=numpy_to_torch_dtype_dict[dtype]).to(
+                device=device
+            )
+            self.tensors[name] = tensor
+
+    def __call__(self, feed_dict, stream, use_cuda_graph=False):
+        for name, buf in feed_dict.items():
+            self.tensors[name].copy_(buf)
+
+        for name, tensor in self.tensors.items():
+            self.context.set_tensor_address(name, tensor.data_ptr())  # type: ignore
+
+        if use_cuda_graph:
+            if self.cuda_graph_instance is not None:
+                cuassert(cudart.cudaGraphLaunch(self.cuda_graph_instance, stream))
+                cuassert(cudart.cudaStreamSynchronize(stream))
+            else:
+                # do inference before CUDA graph capture
+                noerror = self.context.execute_async_v3(stream)  # type: ignore
+                if not noerror:
+                    raise ValueError("ERROR: inference failed.")
+                # capture cuda graph
+                cuassert(
+                    cudart.cudaStreamBeginCapture(
+                        stream, cudart.cudaStreamCaptureMode.cudaStreamCaptureModeGlobal
+                    )
+                )
+                self.context.execute_async_v3(stream)  # type: ignore
+                self.graph = cuassert(cudart.cudaStreamEndCapture(stream))
+                self.cuda_graph_instance = cuassert(cudart.cudaGraphInstantiate(self.graph, 0))
+        else:
+            noerror = self.context.execute_async_v3(stream)  # type: ignore
+            if not noerror:
+                raise ValueError("ERROR: inference failed.")
+
+        return self.tensors
+
+
+def cuassert(cuda_ret):
+    err = cuda_ret[0]
+    if err != cudart.cudaError_t.cudaSuccess:
+        raise RuntimeError(
+            f"CUDA ERROR: {err}, error code reference: https://nvidia.github.io/cuda-python/module/cudart.html#cuda.cudart.cudaError_t"
+        )
+    if len(cuda_ret) > 1:
+        return cuda_ret[1]
+    return None

uv.lock

@@ -34,19 +34,6 @@ version = "4.9.3"
 source = { registry = "https://pypi.org/simple" }
 sdist = { url = "https://files.pythonhosted.org/packages/3e/38/7859ff46355f76f8d19459005ca000b6e7012f2f1ca597746cbcd1fbfe5e/antlr4-python3-runtime-4.9.3.tar.gz", hash = "sha256:f224469b4168294902bb1efa80a8bf7855f24c99aef99cbefc1bcd3cce77881b", size = 117034 }
 
-[[package]]
-name = "bitsandbytes"
-version = "0.44.1"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "numpy" },
-    { name = "torch" },
-]
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/e4/e6/ccb84da7ffaf208a71c2c3c8e1120b34759df640db959660be9a98505eb4/bitsandbytes-0.44.1-py3-none-manylinux_2_24_x86_64.whl", hash = "sha256:b2f24c6cbf11fc8c5d69b3dcecee9f7011451ec59d6ac833e873c9f105259668", size = 122419627 },
-    { url = "https://files.pythonhosted.org/packages/5f/f5/11bddebb5addc0a005b0c1cecc6e4c6e4055ad7b860bdcbf6374e12a51f5/bitsandbytes-0.44.1-py3-none-win_amd64.whl", hash = "sha256:8e68e12aa25d2cf9a1730ad72890a5d1a19daa23f459a6a4679331f353d58cb4", size = 121451331 },
-]
-
 [[package]]
 name = "certifi"
 version = "2024.8.30"
@@ -89,9 +76,19 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/d1/d6/3965ed04c63042e047cb6a3e6ed1a63a35087b6a609aa3a15ed8ac56c221/colorama-0.4.6-py2.py3-none-any.whl", hash = "sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6", size = 25335 },
 ]
 
+[[package]]
+name = "cuda-python"
+version = "12.6.0"
+source = { registry = "https://pypi.org/simple" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/0b/a3/ad3148d068d78e8ad1e40094ab787338ea4bef06fbe2915cf1557a5c5f98/cuda_python-12.6.0-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:dee03e2ba78a807a907a7939dddf089bb8a780faaf7ccbcbfc2461090af11e78", size = 23793330 },
+    { url = "https://files.pythonhosted.org/packages/86/93/f00a5f48eb67216d8a8818b93c0e8bbe5949f297add3367522081ec5223c/cuda_python-12.6.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:e177f584094d9c9fd9c7d153168486a3966765c79cb2a80e86feb15e3b5adc14", size = 24223726 },
+    { url = "https://files.pythonhosted.org/packages/f6/e0/c2302ff6796eac6c6f1e1414f163c6a38deba62af0b7df2b77562656188c/cuda_python-12.6.0-cp310-cp310-win_amd64.whl", hash = "sha256:3b1e9711c6455fabd947076d52eb21ea508ade95eb4dd33838b0339a84238125", size = 9995130 },
+]
+
 [[package]]
 name = "diffusers"
-version = "0.28.2"
+version = "0.30.2"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "filelock" },
@@ -103,9 +100,9 @@ dependencies = [
     { name = "requests" },
     { name = "safetensors" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/45/aa/f40919df7672c278b10509bb0f648a2f4b12ebb82174c0e51a25cbf76421/diffusers-0.28.2.tar.gz", hash = "sha256:0e1e881c821bd6dc2d88252a228a2532ae3577eb4f8b351c9c4abb184aee9b3f", size = 1704805 }
+sdist = { url = "https://files.pythonhosted.org/packages/04/ee/13a6327f04f21420ab4d8ada635aba7d884bf57b09f9b847b9af3818b348/diffusers-0.30.2.tar.gz", hash = "sha256:641875f78f36bdfa4b9af752b124d1fd6d431eadd5547fe0a3f354ae0af2636c", size = 2095560 }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/8c/1d/decefc6459a40df26e17631e42fb7c3814f8424037a26d4bad187800179e/diffusers-0.28.2-py3-none-any.whl", hash = "sha256:9713b7666510723ad6b7e8f4702664ef20c59e4bc241d26c7ca086be3041082e", size = 2184296 },
+    { url = "https://files.pythonhosted.org/packages/2f/ee/f67b0888229be96a276257579a58eb2331733d246fdb8620e09ca7253971/diffusers-0.30.2-py3-none-any.whl", hash = "sha256:739826043147c2b59560944591dfdea5d24cd4fb15e751abbe20679a289bece8", size = 2636928 },
 ]
 
 [[package]]
@@ -114,51 +111,43 @@ version = "6"
 source = { editable = "." }
 dependencies = [
     { name = "accelerate" },
-    { name = "bitsandbytes" },
+    { name = "cuda-python" },
     { name = "diffusers" },
     { name = "edge-maxxing-pipelines" },
-    { name = "huggingface-hub" },
-    { name = "numpy" },
     { name = "omegaconf" },
-    { name = "onediff" },
-    { name = "onediffx" },
-    { name = "oneflow" },
+    { name = "onnx" },
+    { name = "polygraphy" },
     { name = "setuptools" },
-    { name = "stable-fast" },
-    { name = "tomesd" },
+    { name = "tensorrt" },
+    { name = "tensorrt-cu12-bindings" },
+    { name = "tensorrt-cu12-libs" },
     { name = "torch" },
-    { name = "torchvision" },
     { name = "transformers" },
-    { name = "triton" },
-    { name = "xformers" },
+    { name = "wheel" },
 ]
 
 [package.metadata]
 requires-dist = [
     { name = "accelerate", specifier = "==0.31.0" },
-    { name = "bitsandbytes", specifier = ">=0.44.1" },
-    { name = "diffusers", specifier = "==0.28.2" },
-    { name = "edge-maxxing-pipelines", git = "https://github.com/womboai/edge-maxxing?subdirectory=pipelines" },
-    { name = "huggingface-hub", specifier = "==0.25.2" },
-    { name = "numpy", specifier = "==1.26.4" },
+    { name = "cuda-python", specifier = ">=12.6.0" },
+    { name = "diffusers", specifier = "==0.30.2" },
+    { name = "edge-maxxing-pipelines", git = "https://github.com/womboai/edge-maxxing?subdirectory=pipelines&rev=8d8ff45863416484b5b4bc547782591bbdfc696a#8d8ff45863416484b5b4bc547782591bbdfc696a" },
     { name = "omegaconf", specifier = "==2.3.0" },
-    { name = "onediff", specifier = "==1.2.0" },
-    { name = "onediffx", specifier = "==1.2.0" },
-    { name = "oneflow", url = "https://github.com/siliconflow/oneflow_releases/releases/download/community_cu118/oneflow-0.9.1.dev20240802%2Bcu118-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl" },
+    { name = "onnx" },
+    { name = "polygraphy" },
     { name = "setuptools", specifier = ">=75.2.0" },
-    { name = "stable-fast", url = "https://github.com/chengzeyi/stable-fast/releases/download/v1.0.5/stable_fast-1.0.5+torch222cu121-cp310-cp310-manylinux2014_x86_64.whl" },
-    { name = "tomesd", specifier = ">=0.1.3" },
-    { name = "torch", specifier = "==2.2.2" },
-    { name = "torchvision", specifier = "==0.17.2" },
+    { name = "tensorrt", specifier = ">=10.5.0" },
+    { name = "tensorrt-cu12-bindings", specifier = ">=10.5.0" },
+    { name = "tensorrt-cu12-libs", specifier = ">=10.5.0" },
+    { name = "torch", specifier = "==2.4.1" },
     { name = "transformers", specifier = "==4.41.2" },
-    { name = "triton", specifier = "==2.2.0" },
-    { name = "xformers", specifier = "==0.0.25.post1" },
+    { name = "wheel" },
 ]
 
 [[package]]
 name = "edge-maxxing-pipelines"
 version = "1.0.0"
-source = { git = "https://github.com/womboai/edge-maxxing?subdirectory=pipelines#5e655562d2587b575f3b2ad50e8c40560d79fd6f" }
+source = { git = "https://github.com/womboai/edge-maxxing?subdirectory=pipelines&rev=8d8ff45863416484b5b4bc547782591bbdfc696a#8d8ff45863416484b5b4bc547782591bbdfc696a" }
 dependencies = [
     { name = "pydantic" },
 ]
@@ -183,7 +172,7 @@ wheels = [
 
 [[package]]
 name = "huggingface-hub"
-version = "0.25.2"
+version = "0.26.1"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "filelock" },
@@ -194,9 +183,9 @@ dependencies = [
     { name = "tqdm" },
     { name = "typing-extensions" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/df/fd/5f81bae67096c5ab50d29a0230b8374f0245916cca192f8ee2fada51f4f6/huggingface_hub-0.25.2.tar.gz", hash = "sha256:a1014ea111a5f40ccd23f7f7ba8ac46e20fa3b658ced1f86a00c75c06ec6423c", size = 365806 }
+sdist = { url = "https://files.pythonhosted.org/packages/44/99/c8fdef6fe09a1719e5e5de24b012de5824889168c96143f5531cab5af42b/huggingface_hub-0.26.1.tar.gz", hash = "sha256:414c0d9b769eecc86c70f9d939d0f48bb28e8461dd1130021542eff0212db890", size = 375458 }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/64/09/a535946bf2dc88e61341f39dc507530411bb3ea4eac493e5ec833e8f35bd/huggingface_hub-0.25.2-py3-none-any.whl", hash = "sha256:1897caf88ce7f97fe0110603d8f66ac264e3ba6accdf30cd66cc0fed5282ad25", size = 436575 },
+    { url = "https://files.pythonhosted.org/packages/d7/4d/017d8d7cff5100092da8ea19139bcb1965bbadcbb5ddd0480e2badc299e8/huggingface_hub-0.26.1-py3-none-any.whl", hash = "sha256:5927a8fc64ae68859cd954b7cc29d1c8390a5e15caba6d3d349c973be8fdacf3", size = 447439 },
 ]
 
 [[package]]
@@ -232,18 +221,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/31/80/3a54838c3fb461f6fec263ebf3a3a41771bd05190238de3486aae8540c36/jinja2-3.1.4-py3-none-any.whl", hash = "sha256:bc5dd2abb727a5319567b7a813e6a2e7318c39f4f487cfe6c89c6f9c7d25197d", size = 133271 },
 ]
 
-[[package]]
-name = "markdown-it-py"
-version = "3.0.0"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "mdurl" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/38/71/3b932df36c1a044d397a1f92d1cf91ee0a503d91e470cbd670aa66b07ed0/markdown-it-py-3.0.0.tar.gz", hash = "sha256:e3f60a94fa066dc52ec76661e37c851cb232d92f9886b15cb560aaada2df8feb", size = 74596 }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/42/d7/1ec15b46af6af88f19b8e5ffea08fa375d433c998b8a7639e76935c14f1f/markdown_it_py-3.0.0-py3-none-any.whl", hash = "sha256:355216845c60bd96232cd8d8c40e8f9765cc86f46880e43a8fd22dc1a1a8cab1", size = 87528 },
-]
-
 [[package]]
 name = "markupsafe"
 version = "3.0.2"
@@ -262,15 +239,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/44/06/e7175d06dd6e9172d4a69a72592cb3f7a996a9c396eee29082826449bbc3/MarkupSafe-3.0.2-cp310-cp310-win_amd64.whl", hash = "sha256:6af100e168aa82a50e186c82875a5893c5597a0c1ccdb0d8b40240b1f28b969a", size = 15514 },
 ]
 
-[[package]]
-name = "mdurl"
-version = "0.1.2"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/d6/54/cfe61301667036ec958cb99bd3efefba235e65cdeb9c84d24a8293ba1d90/mdurl-0.1.2.tar.gz", hash = "sha256:bb413d29f5eea38f31dd4754dd7377d4465116fb207585f97bf925588687c1ba", size = 8729 }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/b3/38/89ba8ad64ae25be8de66a6d463314cf1eb366222074cfda9ee839c56a4b4/mdurl-0.1.2-py3-none-any.whl", hash = "sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8", size = 9979 },
-]
-
 [[package]]
 name = "mpmath"
 version = "1.3.0"
@@ -291,29 +259,24 @@ wheels = [
 
 [[package]]
 name = "numpy"
-version = "1.26.4"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/65/6e/09db70a523a96d25e115e71cc56a6f9031e7b8cd166c1ac8438307c14058/numpy-1.26.4.tar.gz", hash = "sha256:2a02aba9ed12e4ac4eb3ea9421c420301a0c6460d9830d74a9df87efa4912010", size = 15786129 }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/a7/94/ace0fdea5241a27d13543ee117cbc65868e82213fb31a8eb7fe9ff23f313/numpy-1.26.4-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:9ff0f4f29c51e2803569d7a51c2304de5554655a60c5d776e35b4a41413830d0", size = 20631468 },
-    { url = "https://files.pythonhosted.org/packages/20/f7/b24208eba89f9d1b58c1668bc6c8c4fd472b20c45573cb767f59d49fb0f6/numpy-1.26.4-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:2e4ee3380d6de9c9ec04745830fd9e2eccb3e6cf790d39d7b98ffd19b0dd754a", size = 13966411 },
-    { url = "https://files.pythonhosted.org/packages/fc/a5/4beee6488160798683eed5bdb7eead455892c3b4e1f78d79d8d3f3b084ac/numpy-1.26.4-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:d209d8969599b27ad20994c8e41936ee0964e6da07478d6c35016bc386b66ad4", size = 14219016 },
-    { url = "https://files.pythonhosted.org/packages/4b/d7/ecf66c1cd12dc28b4040b15ab4d17b773b87fa9d29ca16125de01adb36cd/numpy-1.26.4-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ffa75af20b44f8dba823498024771d5ac50620e6915abac414251bd971b4529f", size = 18240889 },
-    { url = "https://files.pythonhosted.org/packages/24/03/6f229fe3187546435c4f6f89f6d26c129d4f5bed40552899fcf1f0bf9e50/numpy-1.26.4-cp310-cp310-musllinux_1_1_aarch64.whl", hash = "sha256:62b8e4b1e28009ef2846b4c7852046736bab361f7aeadeb6a5b89ebec3c7055a", size = 13876746 },
-    { url = "https://files.pythonhosted.org/packages/39/fe/39ada9b094f01f5a35486577c848fe274e374bbf8d8f472e1423a0bbd26d/numpy-1.26.4-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:a4abb4f9001ad2858e7ac189089c42178fcce737e4169dc61321660f1a96c7d2", size = 18078620 },
-    { url = "https://files.pythonhosted.org/packages/d5/ef/6ad11d51197aad206a9ad2286dc1aac6a378059e06e8cf22cd08ed4f20dc/numpy-1.26.4-cp310-cp310-win32.whl", hash = "sha256:bfe25acf8b437eb2a8b2d49d443800a5f18508cd811fea3181723922a8a82b07", size = 5972659 },
-    { url = "https://files.pythonhosted.org/packages/19/77/538f202862b9183f54108557bfda67e17603fc560c384559e769321c9d92/numpy-1.26.4-cp310-cp310-win_amd64.whl", hash = "sha256:b97fe8060236edf3662adfc2c633f56a08ae30560c56310562cb4f95500022d5", size = 15808905 },
-]
-
-[[package]]
-name = "nvidia-cublas-cu11"
-version = "11.11.3.6"
+version = "2.1.2"
 source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/4b/d1/8a730ea07f4a37d94f9172f4ce1d81064b7a64766b460378be278952de75/numpy-2.1.2.tar.gz", hash = "sha256:13532a088217fa624c99b843eeb54640de23b3414b14aa66d023805eb731066c", size = 18878063 }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/46/be/c222e33e60d28ecd496a46fc4d78ccae0ee28e1fd7dc705b6288b4cad27e/nvidia_cublas_cu11-11.11.3.6-py3-none-manylinux1_x86_64.whl", hash = "sha256:39fb40e8f486dd8a2ddb8fdeefe1d5b28f5b99df01c87ab3676f057a74a5a6f3", size = 417870452 },
-    { url = "https://files.pythonhosted.org/packages/96/df/c5ac9ac5096355c47c606e613ecc7aa50fbccf5e0145df857d11da6464b1/nvidia_cublas_cu11-11.11.3.6-py3-none-manylinux2014_aarch64.whl", hash = "sha256:5ccae9e069a2c6be9af9cb5a0b0c6928c19c7915e390d15f598a1eead2a01a7a", size = 291428448 },
-    { url = "https://files.pythonhosted.org/packages/ea/2e/9d99c60771d275ecf6c914a612e9a577f740a615bc826bec132368e1d3ae/nvidia_cublas_cu11-11.11.3.6-py3-none-manylinux2014_x86_64.whl", hash = "sha256:60252822adea5d0b10cd990a7dc7bedf7435f30ae40083c7a624a85a43225abc", size = 417870460 },
-    { url = "https://files.pythonhosted.org/packages/0b/1d/7a78cd36fd5e3da4021b3ac2c2c8b2651dd72345b7c3ecc0d3e051884f50/nvidia_cublas_cu11-11.11.3.6-py3-none-win_amd64.whl", hash = "sha256:6ab12b1302bef8ac1ff4414edd1c059e57f4833abef9151683fb8f4de25900be", size = 427234740 },
+    { url = "https://files.pythonhosted.org/packages/1c/a2/40a76d357f168e9f9f06d6cc2c8e22dd5fb2bfbe63fe2c433057258c145a/numpy-2.1.2-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:30d53720b726ec36a7f88dc873f0eec8447fbc93d93a8f079dfac2629598d6ee", size = 21150947 },
+    { url = "https://files.pythonhosted.org/packages/b5/d0/ba271ea9108d7278d3889a7eb38d77370a88713fb94339964e71ac184d4a/numpy-2.1.2-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:e8d3ca0a72dd8846eb6f7dfe8f19088060fcb76931ed592d29128e0219652884", size = 13758184 },
+    { url = "https://files.pythonhosted.org/packages/7c/b9/5c6507439cd756201010f7937bf90712c2469052ae094584af14557dd64f/numpy-2.1.2-cp310-cp310-macosx_14_0_arm64.whl", hash = "sha256:fc44e3c68ff00fd991b59092a54350e6e4911152682b4782f68070985aa9e648", size = 5354091 },
+    { url = "https://files.pythonhosted.org/packages/60/21/7938cf724d9e84e45fb886f3fc794ab431d71facfebc261e3e9f19f3233a/numpy-2.1.2-cp310-cp310-macosx_14_0_x86_64.whl", hash = "sha256:7c1c60328bd964b53f8b835df69ae8198659e2b9302ff9ebb7de4e5a5994db3d", size = 6887169 },
+    { url = "https://files.pythonhosted.org/packages/09/8d/42a124657f5d31902fca73921b25a0d022cead2b32ce7e6975762cd2995a/numpy-2.1.2-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:6cdb606a7478f9ad91c6283e238544451e3a95f30fb5467fbf715964341a8a86", size = 13888165 },
+    { url = "https://files.pythonhosted.org/packages/fb/25/ba023652a39a2c127200e85aed975fc6119b421e2c348e5d0171e2046edb/numpy-2.1.2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d666cb72687559689e9906197e3bec7b736764df6a2e58ee265e360663e9baf7", size = 16326954 },
+    { url = "https://files.pythonhosted.org/packages/34/58/23e6b07fad492b7c47cf09cd8bad6983658f0f925b6c535fd008e3e86274/numpy-2.1.2-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:c6eef7a2dbd0abfb0d9eaf78b73017dbfd0b54051102ff4e6a7b2980d5ac1a03", size = 16702916 },
+    { url = "https://files.pythonhosted.org/packages/91/24/37b5cf2dc7d385ac97f7b7fe50cba312abb70a2a5eac74c23af028811f73/numpy-2.1.2-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:12edb90831ff481f7ef5f6bc6431a9d74dc0e5ff401559a71e5e4611d4f2d466", size = 14384372 },
+    { url = "https://files.pythonhosted.org/packages/ea/ec/0f6d471058a01d1a05a50d2793898de1549280fa715a8537987ee866b5d9/numpy-2.1.2-cp310-cp310-win32.whl", hash = "sha256:a65acfdb9c6ebb8368490dbafe83c03c7e277b37e6857f0caeadbbc56e12f4fb", size = 6535361 },
+    { url = "https://files.pythonhosted.org/packages/c2/3d/293cc5927f916a7bc6bf74da8f6defab63d1b13f0959d7e21878ad8a20d8/numpy-2.1.2-cp310-cp310-win_amd64.whl", hash = "sha256:860ec6e63e2c5c2ee5e9121808145c7bf86c96cca9ad396c0bd3e0f2798ccbe2", size = 12865501 },
+    { url = "https://files.pythonhosted.org/packages/73/c9/3e1d6bbe6d3d2e2c5a9483b24b2f29a229b323f62054278a3bba7fee11e5/numpy-2.1.2-pp310-pypy310_pp73-macosx_10_15_x86_64.whl", hash = "sha256:bdd407c40483463898b84490770199d5714dcc9dd9b792f6c6caccc523c00952", size = 20981945 },
+    { url = "https://files.pythonhosted.org/packages/6e/62/989c4988bde1a8e08117fccc3bab73d2886421fb98cde597168714f3c54e/numpy-2.1.2-pp310-pypy310_pp73-macosx_14_0_x86_64.whl", hash = "sha256:da65fb46d4cbb75cb417cddf6ba5e7582eb7bb0b47db4b99c9fe5787ce5d91f5", size = 6750558 },
+    { url = "https://files.pythonhosted.org/packages/53/b1/00ef9f30975f1312a53257f68e57b4513d14d537e03d507e2773a684b1e8/numpy-2.1.2-pp310-pypy310_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:1c193d0b0238638e6fc5f10f1b074a6993cb13b0b431f64079a509d63d3aa8b7", size = 16141552 },
+    { url = "https://files.pythonhosted.org/packages/c0/ec/0c04903b48dfea6be1d7b47ba70f98709fb7198fd970784a1400c391d522/numpy-2.1.2-pp310-pypy310_pp73-win_amd64.whl", hash = "sha256:a7d80b2e904faa63068ead63107189164ca443b42dd1930299e0d1cb041cec2e", size = 12789924 },
 ]
 
 [[package]]
@@ -346,40 +309,18 @@ version = "12.1.105"
 source = { registry = "https://pypi.org/simple" }
 wheels = [
     { url = "https://files.pythonhosted.org/packages/eb/d5/c68b1d2cdfcc59e72e8a5949a37ddb22ae6cade80cd4a57a84d4c8b55472/nvidia_cuda_runtime_cu12-12.1.105-py3-none-manylinux1_x86_64.whl", hash = "sha256:6e258468ddf5796e25f1dc591a31029fa317d97a0a94ed93468fc86301d61e40", size = 823596 },
-]
-
-[[package]]
-name = "nvidia-cudnn-cu11"
-version = "9.5.0.50"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "nvidia-cublas-cu11" },
-]
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/a6/d6/ec0bf8fd29c907e68de69248e5dbc3e78b63a613163d54ebfc4ca9362421/nvidia_cudnn_cu11-9.5.0.50-py3-none-manylinux2014_x86_64.whl", hash = "sha256:15ddd1c8e2904f35debb4a0452f4df57eeb9d363e760de5a5c2ca0daba8e545a", size = 558442361 },
-    { url = "https://files.pythonhosted.org/packages/2e/01/e7aeb8f9f2336a44c946e629eeb2d204ec672ecd94a47d3b1abaa10a05f1/nvidia_cudnn_cu11-9.5.0.50-py3-none-win_amd64.whl", hash = "sha256:ecf46a9ca869c1628f0ce91c160d53cbd28f0805dca856f85f3e65b4ef3496b1", size = 553719005 },
+    { url = "https://files.pythonhosted.org/packages/9f/e2/7a2b4b5064af56ea8ea2d8b2776c0f2960d95c88716138806121ae52a9c9/nvidia_cuda_runtime_cu12-12.1.105-py3-none-win_amd64.whl", hash = "sha256:dfb46ef84d73fababab44cf03e3b83f80700d27ca300e537f85f636fac474344", size = 821226 },
 ]
 
 [[package]]
 name = "nvidia-cudnn-cu12"
-version = "8.9.2.26"
+version = "9.1.0.70"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "nvidia-cublas-cu12" },
 ]
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/ff/74/a2e2be7fb83aaedec84f391f082cf765dfb635e7caa9b49065f73e4835d8/nvidia_cudnn_cu12-8.9.2.26-py3-none-manylinux1_x86_64.whl", hash = "sha256:5ccb288774fdfb07a7e7025ffec286971c06d8d7b4fb162525334616d7629ff9", size = 731725872 },
-]
-
-[[package]]
-name = "nvidia-cufft-cu11"
-version = "10.9.0.58"
-source = { registry = "https://pypi.org/simple" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/74/79/b912a77e38e41f15a0581a59f5c3548d1ddfdda3225936fb67c342719e7a/nvidia_cufft_cu11-10.9.0.58-py3-none-manylinux1_x86_64.whl", hash = "sha256:222f9da70c80384632fd6035e4c3f16762d64ea7a843829cb278f98b3cb7dd81", size = 168405414 },
-    { url = "https://files.pythonhosted.org/packages/71/7a/a2ad9951d57c3cc23f4fa6d84b146afd9f375ffbc744b38935930ac4393f/nvidia_cufft_cu11-10.9.0.58-py3-none-manylinux2014_aarch64.whl", hash = "sha256:34b7315104e615b230dc3c2d1861f13bff9ec465c5d3b4bb65b4986d03a1d8d4", size = 111231060 },
-    { url = "https://files.pythonhosted.org/packages/64/c8/133717b43182ba063803e983e7680a94826a9f4ff5734af0ca315803f1b3/nvidia_cufft_cu11-10.9.0.58-py3-none-manylinux2014_x86_64.whl", hash = "sha256:e21037259995243cc370dd63c430d77ae9280bedb68d5b5a18226bfc92e5d748", size = 168405419 },
-    { url = "https://files.pythonhosted.org/packages/f8/b4/e432a74f8db0e84f734dc14d36c0e529225132bf7e239da21f55893351a6/nvidia_cufft_cu11-10.9.0.58-py3-none-win_amd64.whl", hash = "sha256:c4d316f17c745ec9c728e30409612eaf77a8404c3733cdf6c9c1569634d1ca03", size = 172237004 },
+    { url = "https://files.pythonhosted.org/packages/9f/fd/713452cd72343f682b1c7b9321e23829f00b842ceaedcda96e742ea0b0b3/nvidia_cudnn_cu12-9.1.0.70-py3-none-manylinux2014_x86_64.whl", hash = "sha256:165764f44ef8c61fcdfdfdbe769d687e06374059fbb388b6c89ecb0e28793a6f", size = 664752741 },
 ]
 
 [[package]]
@@ -411,17 +352,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/bc/1d/8de1e5c67099015c834315e333911273a8c6aaba78923dd1d1e25fc5f217/nvidia_cusolver_cu12-11.4.5.107-py3-none-manylinux1_x86_64.whl", hash = "sha256:8a7ec542f0412294b15072fa7dab71d31334014a69f953004ea7a118206fe0dd", size = 124161928 },
 ]
 
-[[package]]
-name = "nvidia-cusparse-cu11"
-version = "11.7.5.86"
-source = { registry = "https://pypi.org/simple" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/c1/e0/21b829c535d569831835a4ca5d049a19ba00d3e91f3e12ab4ad27bd7385f/nvidia_cusparse_cu11-11.7.5.86-py3-none-manylinux1_x86_64.whl", hash = "sha256:4ae709fe78d3f23f60acaba8c54b8ad556cf16ca486e0cc1aa92dca7555d2d2b", size = 204126221 },
-    { url = "https://files.pythonhosted.org/packages/a2/6e/4eb2842e7ab1804072bca43030a8b92731e5a35f6a4a1b8f1aa8fa5f411c/nvidia_cusparse_cu11-11.7.5.86-py3-none-manylinux2014_aarch64.whl", hash = "sha256:6c7da46abee7567e619d4aa2e90a1b032cfcbd1211d429853b1a6e87514a14b2", size = 203917797 },
-    { url = "https://files.pythonhosted.org/packages/ed/5c/b0333b07c51ced77397c2fb0d9826072cea0da9d421aa7e792aa0f8ecc72/nvidia_cusparse_cu11-11.7.5.86-py3-none-manylinux2014_x86_64.whl", hash = "sha256:8d7cf1628fd8d462b5d2ba6678fae34733a48ecb80495b9c68672ec6a6dde5ef", size = 204126227 },
-    { url = "https://files.pythonhosted.org/packages/b8/36/a670e8ca1deccd3c63be4d0286491cf5c6375253f0d948e7cc5167fe1da9/nvidia_cusparse_cu11-11.7.5.86-py3-none-win_amd64.whl", hash = "sha256:a0f6ee81cd91be606fc2f55992d06b09cd4e86d74b6ae5e8dd1631cf7f5a8706", size = 203420667 },
-]
-
 [[package]]
 name = "nvidia-cusparse-cu12"
 version = "12.1.0.106"
@@ -433,20 +363,13 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/65/5b/cfaeebf25cd9fdec14338ccb16f6b2c4c7fa9163aefcf057d86b9cc248bb/nvidia_cusparse_cu12-12.1.0.106-py3-none-manylinux1_x86_64.whl", hash = "sha256:f3b50f42cf363f86ab21f720998517a659a48131e8d538dc02f8768237bd884c", size = 195958278 },
 ]
 
-[[package]]
-name = "nvidia-nccl-cu11"
-version = "2.21.5"
-source = { registry = "https://pypi.org/simple" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/ac/9a/8b6a28b3b87d5fddab0e92cd835339eb8fbddaa71ae67518c8c1b3d05bae/nvidia_nccl_cu11-2.21.5-py3-none-manylinux2014_x86_64.whl", hash = "sha256:49d8350629c7888701d1fd200934942671cb5c728f49acc5a0b3a768820bed29", size = 147811630 },
-]
-
 [[package]]
 name = "nvidia-nccl-cu12"
-version = "2.19.3"
+version = "2.20.5"
 source = { registry = "https://pypi.org/simple" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/38/00/d0d4e48aef772ad5aebcf70b73028f88db6e5640b36c38e90445b7a57c45/nvidia_nccl_cu12-2.19.3-py3-none-manylinux1_x86_64.whl", hash = "sha256:a9734707a2c96443331c1e48c717024aa6678a0e2a4cb66b2c364d18cee6b48d", size = 165987969 },
+    { url = "https://files.pythonhosted.org/packages/c1/bb/d09dda47c881f9ff504afd6f9ca4f502ded6d8fc2f572cacc5e39da91c28/nvidia_nccl_cu12-2.20.5-py3-none-manylinux2014_aarch64.whl", hash = "sha256:1fc150d5c3250b170b29410ba682384b14581db722b2531b0d8d33c595f33d01", size = 176238458 },
+    { url = "https://files.pythonhosted.org/packages/4b/2a/0a131f572aa09f741c30ccd45a8e56316e8be8dfc7bc19bf0ab7cfef7b19/nvidia_nccl_cu12-2.20.5-py3-none-manylinux2014_x86_64.whl", hash = "sha256:057f6bf9685f75215d0c53bf3ac4a10b3e6578351de307abad9e18a99182af56", size = 176249402 },
 ]
 
 [[package]]
@@ -480,79 +403,22 @@ wheels = [
 ]
 
 [[package]]
-name = "onediff"
-version = "1.2.0"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "accelerate" },
-    { name = "diffusers" },
-    { name = "torch" },
-    { name = "transformers" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/d5/30/b493cbca73c5cdd2499a3ec37e451fa4c826e02be6bb6bc66c44d98293cc/onediff-1.2.0.tar.gz", hash = "sha256:8655b18698ca252093c06db53d4620de4065d1e507f9d7b1cbe3f123ed17bd26", size = 73823417 }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/3c/22/5b6de9dae4b8ed161bfcab9a7d6d7660cca4ee1d7db733950b67905f7f66/onediff-1.2.0-py3-none-any.whl", hash = "sha256:84bb2285ae7f38e450229779ab96b591e4f530e1f22b7320f41b06f7aec181a3", size = 104932 },
-]
-
-[[package]]
-name = "onediffx"
-version = "1.2.0"
+name = "onnx"
+version = "1.17.0"
 source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "accelerate" },
-    { name = "diffusers" },
-    { name = "omegaconf" },
-    { name = "onefx" },
-    { name = "torch" },
-    { name = "transformers" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/b9/e1/b0797217862a9eb48bdb5f3e53ed8c9b69fe582de1e7c3a4b948f7ed862a/onediffx-1.2.0.tar.gz", hash = "sha256:3e0934e71f966d41d04951826ae36ebfb25be8c23ca7f5e6f7b177fd92f0791b", size = 95945 }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/f4/dd/412be6d3646355b5e5b75066f715fc34d0bbdcfbf18d0934b9f107c59239/onediffx-1.2.0-py3-none-any.whl", hash = "sha256:86dea026ac30bc04a2f4cd25ac4a511b1c4b4154657bc48c7e05f4a238085f14", size = 67899 },
-]
-
-[[package]]
-name = "oneflow"
-version = "0.9.1.dev20240802+cu118"
-source = { url = "https://github.com/siliconflow/oneflow_releases/releases/download/community_cu118/oneflow-0.9.1.dev20240802%2Bcu118-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl" }
 dependencies = [
     { name = "numpy" },
-    { name = "nvidia-cublas-cu11" },
-    { name = "nvidia-cudnn-cu11" },
-    { name = "nvidia-cufft-cu11" },
-    { name = "nvidia-cusparse-cu11" },
-    { name = "nvidia-nccl-cu11" },
-    { name = "pillow" },
     { name = "protobuf" },
-    { name = "requests" },
-    { name = "rich" },
-    { name = "tqdm" },
 ]
+sdist = { url = "https://files.pythonhosted.org/packages/9a/54/0e385c26bf230d223810a9c7d06628d954008a5e5e4b73ee26ef02327282/onnx-1.17.0.tar.gz", hash = "sha256:48ca1a91ff73c1d5e3ea2eef20ae5d0e709bb8a2355ed798ffc2169753013fd3", size = 12165120 }
 wheels = [
-    { url = "https://github.com/siliconflow/oneflow_releases/releases/download/community_cu118/oneflow-0.9.1.dev20240802%2Bcu118-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:0a9de80a3c059ec78bec7dfa6952f5618aaa05f29182b49634e0211195281e19" },
-]
-
-[package.metadata]
-requires-dist = [
-    { name = "numpy", specifier = ">=1.22.1" },
-    { name = "nvidia-cublas-cu11" },
-    { name = "nvidia-cudnn-cu11" },
-    { name = "nvidia-cufft-cu11" },
-    { name = "nvidia-cusparse-cu11" },
-    { name = "nvidia-nccl-cu11" },
-    { name = "pillow" },
-    { name = "protobuf", specifier = ">=3.9.2" },
-    { name = "requests" },
-    { name = "rich" },
-    { name = "tqdm" },
+    { url = "https://files.pythonhosted.org/packages/2e/29/57053ba7787788ac75efb095cfc1ae290436b6d3a26754693cd7ed1b4fac/onnx-1.17.0-cp310-cp310-macosx_12_0_universal2.whl", hash = "sha256:38b5df0eb22012198cdcee527cc5f917f09cce1f88a69248aaca22bd78a7f023", size = 16645616 },
+    { url = "https://files.pythonhosted.org/packages/75/0d/831807a18db2a5e8f7813848c59272b904a4ef3939fe4d1288cbce9ea735/onnx-1.17.0-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:d545335cb49d4d8c47cc803d3a805deb7ad5d9094dc67657d66e568610a36d7d", size = 15908420 },
+    { url = "https://files.pythonhosted.org/packages/dd/5b/c4f95dbe652d14aeba9afaceb177e9ffc48ac3c03048dd3f872f26f07e34/onnx-1.17.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:3193a3672fc60f1a18c0f4c93ac81b761bc72fd8a6c2035fa79ff5969f07713e", size = 16046244 },
+    { url = "https://files.pythonhosted.org/packages/08/a9/c1f218085043dccc6311460239e253fa6957cf12ee4b0a56b82014938d0b/onnx-1.17.0-cp310-cp310-win32.whl", hash = "sha256:0141c2ce806c474b667b7e4499164227ef594584da432fd5613ec17c1855e311", size = 14423516 },
+    { url = "https://files.pythonhosted.org/packages/0e/d3/d26ebf590a65686dde6b27fef32493026c5be9e42083340d947395f93405/onnx-1.17.0-cp310-cp310-win_amd64.whl", hash = "sha256:dfd777d95c158437fda6b34758f0877d15b89cbe9ff45affbedc519b35345cf9", size = 14528496 },
 ]
 
-[[package]]
-name = "onefx"
-version = "0.0.3"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/3c/c4/2dc5fd8ab613d32e24d54204b4d448b315fa50e5f1ce3197938f2219fa09/onefx-0.0.3.tar.gz", hash = "sha256:d3f3f816891cdfdb6d69e4a78b72d5cafd7d40e165ce419f89f855f47a3dbe34", size = 147260 }
-
 [[package]]
 name = "packaging"
 version = "24.1"
@@ -588,6 +454,14 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/ec/3d/c32a51d848401bd94cabb8767a39621496491ee7cd5199856b77da9b18ad/pillow-11.0.0-pp310-pypy310_pp73-win_amd64.whl", hash = "sha256:224aaa38177597bb179f3ec87eeefcce8e4f85e608025e9cfac60de237ba6316", size = 2567508 },
 ]
 
+[[package]]
+name = "polygraphy"
+version = "0.49.9"
+source = { registry = "https://pypi.org/simple" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/4a/f5/a2b20c677c1a856cc9e08cd0b5a5105450ed5253e369e938ddd31d91c547/polygraphy-0.49.9-py2.py3-none-any.whl", hash = "sha256:62ae22825efdd3288222e5b1d2d791fe58e87844fcd848bcd1251fbce02ba956", size = 346910 },
+]
+
 [[package]]
 name = "protobuf"
 version = "5.28.3"
@@ -662,15 +536,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/a9/f9/b6bcaf874f410564a78908739c80861a171788ef4d4f76f5009656672dfe/pydantic_core-2.23.4-pp310-pypy310_pp73-win_amd64.whl", hash = "sha256:9a5bce9d23aac8f0cf0836ecfc033896aa8443b501c58d0602dbfd5bd5b37753", size = 1920344 },
 ]
 
-[[package]]
-name = "pygments"
-version = "2.18.0"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/8e/62/8336eff65bcbc8e4cb5d05b55faf041285951b6e80f33e2bff2024788f31/pygments-2.18.0.tar.gz", hash = "sha256:786ff802f32e91311bff3889f6e9a86e81505fe99f2735bb6d60ae0c5004f199", size = 4891905 }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/f7/3f/01c8b82017c199075f8f788d0d906b9ffbbc5a47dc9918a945e13d5a2bda/pygments-2.18.0-py3-none-any.whl", hash = "sha256:b8e6aca0523f3ab76fee51799c488e38782ac06eafcf95e7ba832985c8e7b13a", size = 1205513 },
-]
-
 [[package]]
 name = "pyyaml"
 version = "6.0.2"
@@ -727,20 +592,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/f9/9b/335f9764261e915ed497fcdeb11df5dfd6f7bf257d4a6a2a686d80da4d54/requests-2.32.3-py3-none-any.whl", hash = "sha256:70761cfe03c773ceb22aa2f671b4757976145175cdfca038c02654d061d6dcc6", size = 64928 },
 ]
 
-[[package]]
-name = "rich"
-version = "13.9.3"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "markdown-it-py" },
-    { name = "pygments" },
-    { name = "typing-extensions" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/d9/e9/cf9ef5245d835065e6673781dbd4b8911d352fb770d56cf0879cf11b7ee1/rich-13.9.3.tar.gz", hash = "sha256:bc1e01b899537598cf02579d2b9f4a415104d3fc439313a7a2c165d76557a08e", size = 222889 }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/9a/e2/10e9819cf4a20bd8ea2f5dabafc2e6bf4a78d6a0965daeb60a4b34d1c11f/rich-13.9.3-py3-none-any.whl", hash = "sha256:9836f5096eb2172c9e77df411c1b009bace4193d6a481d534fea75ebba758283", size = 242157 },
-]
-
 [[package]]
 name = "safetensors"
 version = "0.4.5"
@@ -778,41 +629,50 @@ wheels = [
 ]
 
 [[package]]
-name = "stable-fast"
-version = "1.0.5+torch222cu121"
-source = { url = "https://github.com/chengzeyi/stable-fast/releases/download/v1.0.5/stable_fast-1.0.5+torch222cu121-cp310-cp310-manylinux2014_x86_64.whl" }
+name = "sympy"
+version = "1.13.3"
+source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "torch" },
+    { name = "mpmath" },
 ]
+sdist = { url = "https://files.pythonhosted.org/packages/11/8a/5a7fd6284fa8caac23a26c9ddf9c30485a48169344b4bd3b0f02fef1890f/sympy-1.13.3.tar.gz", hash = "sha256:b27fd2c6530e0ab39e275fc9b683895367e51d5da91baa8d3d64db2565fec4d9", size = 7533196 }
 wheels = [
-    { url = "https://github.com/chengzeyi/stable-fast/releases/download/v1.0.5/stable_fast-1.0.5+torch222cu121-cp310-cp310-manylinux2014_x86_64.whl", hash = "sha256:03d193666c52016096ecfdf0e8e2183fb5aa8ea51e99d5132b353a1e9a6c1264" },
-]
-
-[package.metadata]
-requires-dist = [
-    { name = "diffusers", marker = "extra == 'diffusers'", specifier = ">=0.19.0" },
-    { name = "numpy", marker = "extra == 'dev'" },
-    { name = "opencv-python", marker = "extra == 'dev'" },
-    { name = "pillow", marker = "extra == 'dev'" },
-    { name = "prettytable", marker = "extra == 'dev'" },
-    { name = "pytest", marker = "extra == 'dev'" },
-    { name = "torch" },
-    { name = "transformers", marker = "extra == 'diffusers'" },
-    { name = "triton", marker = "extra == 'triton'", specifier = ">=2.1.0" },
-    { name = "xformers", marker = "extra == 'xformers'", specifier = ">=0.0.20" },
+    { url = "https://files.pythonhosted.org/packages/99/ff/c87e0622b1dadea79d2fb0b25ade9ed98954c9033722eb707053d310d4f3/sympy-1.13.3-py3-none-any.whl", hash = "sha256:54612cf55a62755ee71824ce692986f23c88ffa77207b30c1368eda4a7060f73", size = 6189483 },
 ]
 
 [[package]]
-name = "sympy"
-version = "1.13.3"
+name = "tensorrt"
+version = "10.5.0"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "mpmath" },
+    { name = "tensorrt-cu12" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/11/8a/5a7fd6284fa8caac23a26c9ddf9c30485a48169344b4bd3b0f02fef1890f/sympy-1.13.3.tar.gz", hash = "sha256:b27fd2c6530e0ab39e275fc9b683895367e51d5da91baa8d3d64db2565fec4d9", size = 7533196 }
+sdist = { url = "https://files.pythonhosted.org/packages/ee/b9/f917eb7dfe02da30bc91206a464c850f4b94a1e14b8f95870074c9b9abea/tensorrt-10.5.0.tar.gz", hash = "sha256:d5c6338d44aeda20250fdbe31f9df8ca152b830f811aaf19d6c4d1dafd18c84b", size = 16401 }
+
+[[package]]
+name = "tensorrt-cu12"
+version = "10.5.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/22/d5/a4c3e22482d4273e151123990934d7c8d0ba1e4efb9a483eba807cdce279/tensorrt-cu12-10.5.0.tar.gz", hash = "sha256:46edbda08c54c8ffa88c75d75b4761eb9839e81678135e8d1530adc8cef6a61b", size = 18341 }
+
+[[package]]
+name = "tensorrt-cu12-bindings"
+version = "10.5.0"
+source = { registry = "https://pypi.org/simple" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/99/ff/c87e0622b1dadea79d2fb0b25ade9ed98954c9033722eb707053d310d4f3/sympy-1.13.3-py3-none-any.whl", hash = "sha256:54612cf55a62755ee71824ce692986f23c88ffa77207b30c1368eda4a7060f73", size = 6189483 },
+    { url = "https://files.pythonhosted.org/packages/21/be/cab39a2c387887fa87bb8f199d113a10ebd0ba8b052927c2ae43b1495cf6/tensorrt_cu12_bindings-10.5.0-cp310-none-manylinux_2_17_x86_64.whl", hash = "sha256:45a31cc3f25489bb05fc9cb8dae0e63b205bf3da1656c44430f97cf263d5720c", size = 1117215 },
+    { url = "https://files.pythonhosted.org/packages/02/49/36db3b3c0bd0c7dc68964c75b1691b46abe8388708b4da04c3261f8ab7c0/tensorrt_cu12_bindings-10.5.0-cp310-none-manylinux_2_31_aarch64.whl", hash = "sha256:900b87824ebbc9e1059a4a9a5ed3040eb9d74ba9a601674086030d373996692a", size = 1091646 },
+    { url = "https://files.pythonhosted.org/packages/17/df/e95a92fa4d43df918cc8bc681697b1423a988db339af25bafe25068c522d/tensorrt_cu12_bindings-10.5.0-cp310-none-win_amd64.whl", hash = "sha256:2bf2eb6d36ed9fe44a4b416def538775012abec34fdb5a6fb8461dd569717055", size = 769305 },
+]
+
+[[package]]
+name = "tensorrt-cu12-libs"
+version = "10.5.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "nvidia-cuda-runtime-cu12" },
 ]
+sdist = { url = "https://files.pythonhosted.org/packages/ff/d2/28d4bdadcb4690e7c051ae23ac5559dffca7ee6bf859ea76c9ab9931ba53/tensorrt_cu12_libs-10.5.0.tar.gz", hash = "sha256:358b3a36c30ab74ad710f227b410206ae94e8d1003c09b75216e39813dac0d9d", size = 630 }
 
 [[package]]
 name = "tokenizers"
@@ -844,21 +704,9 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/45/b6/36c1bb106bbe96012c9367df89ed01599cada036c0b96d38fbbdbeb75c9f/tokenizers-0.19.1-pp310-pypy310_pp73-musllinux_1_1_x86_64.whl", hash = "sha256:43350270bfc16b06ad3f6f07eab21f089adb835544417afda0f83256a8bf8b75", size = 9945103 },
 ]
 
-[[package]]
-name = "tomesd"
-version = "0.1.3"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "torch" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/29/37/ed74c7449fe5a8a4726be3dab4d879d000babf6ea538658171933b1f000e/tomesd-0.1.3.tar.gz", hash = "sha256:15bba2e952f4643c8355951e892fda918ddccbdff2238dc368d42bd078fcedc9", size = 14032 }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/0c/02/367c67c8f510313f143a7818e92254a5f861c7d94c98ad6a08d25db52fee/tomesd-0.1.3-py3-none-any.whl", hash = "sha256:3d5aa0857fe2c2aab253891050601ca13a87d8d7a99b6760b9ca0856aa0c6355", size = 11467 },
-]
-
 [[package]]
 name = "torch"
-version = "2.2.2"
+version = "2.4.1"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "filelock" },
@@ -881,28 +729,10 @@ dependencies = [
     { name = "typing-extensions" },
 ]
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/33/b3/1fcc3bccfddadfd6845dcbfe26eb4b099f1dfea5aa0e5cfb92b3c98dba5b/torch-2.2.2-cp310-cp310-manylinux1_x86_64.whl", hash = "sha256:bc889d311a855dd2dfd164daf8cc903a6b7273a747189cebafdd89106e4ad585", size = 755526581 },
-    { url = "https://files.pythonhosted.org/packages/c3/7c/aeb0c5789a3f10cf909640530cd75b314959b9d9914a4996ed2c7bf8779d/torch-2.2.2-cp310-cp310-manylinux2014_aarch64.whl", hash = "sha256:15dffa4cc3261fa73d02f0ed25f5fa49ecc9e12bf1ae0a4c1e7a88bbfaad9030", size = 86623646 },
-    { url = "https://files.pythonhosted.org/packages/3a/81/684d99e536b20e869a7c1222cf1dd233311fb05d3628e9570992bfb65760/torch-2.2.2-cp310-cp310-win_amd64.whl", hash = "sha256:11e8fe261233aeabd67696d6b993eeb0896faa175c6b41b9a6c9f0334bdad1c5", size = 198579616 },
-    { url = "https://files.pythonhosted.org/packages/3b/55/7192974ab13e5e5577f45d14ce70d42f5a9a686b4f57bbe8c9ab45c4a61a/torch-2.2.2-cp310-none-macosx_10_9_x86_64.whl", hash = "sha256:b2e2200b245bd9f263a0d41b6a2dab69c4aca635a01b30cca78064b0ef5b109e", size = 150788930 },
-    { url = "https://files.pythonhosted.org/packages/33/6b/21496316c9b8242749ee2a9064406271efdf979e91d440e8a3806b5e84bf/torch-2.2.2-cp310-none-macosx_11_0_arm64.whl", hash = "sha256:877b3e6593b5e00b35bbe111b7057464e76a7dd186a287280d941b564b0563c2", size = 59707286 },
-]
-
-[[package]]
-name = "torchvision"
-version = "0.17.2"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "numpy" },
-    { name = "pillow" },
-    { name = "torch" },
-]
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/a2/70/c781e0433ab7b8f6f693580e5065ae17c3785b2836200311765f99535ef8/torchvision-0.17.2-cp310-cp310-macosx_10_13_x86_64.whl", hash = "sha256:1f2910fe3c21ad6875b2720d46fad835b2e4b336e9553d31ca364d24c90b1d4f", size = 1666426 },
-    { url = "https://files.pythonhosted.org/packages/64/3d/a0385fd301e6e2eefb1d9cf2ac97d0c33ebf4a764aba0f066e6e16324916/torchvision-0.17.2-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:ecc1c503fa8a54fbab777e06a7c228032b8ab78efebf35b28bc8f22f544f51f1", size = 1571154 },
-    { url = "https://files.pythonhosted.org/packages/e0/2f/d13cb0ffc4808f85b880ef66ab6cfef10bd35e5c151dae68ea18cf6bf636/torchvision-0.17.2-cp310-cp310-manylinux1_x86_64.whl", hash = "sha256:f400145fc108833e7c2fc28486a04989ca742146d7a2a2cc48878ebbb40cdbbd", size = 6915896 },
-    { url = "https://files.pythonhosted.org/packages/fb/41/c8a440ebb1d4958baf02f08f6df56a3a30bea3fdcdc99076cb7da023babe/torchvision-0.17.2-cp310-cp310-manylinux2014_aarch64.whl", hash = "sha256:e9e4bed404af33dfc92eecc2b513d21ddc4c242a7fd8708b3b09d3a26aa6f444", size = 14008514 },
-    { url = "https://files.pythonhosted.org/packages/0a/b8/027b3c36e61a26d14d4e9a8413d4a09e5fd8d3e01e3efce78447ca1dc3dd/torchvision-0.17.2-cp310-cp310-win_amd64.whl", hash = "sha256:ba2e62f233eab3d42b648c122a3a29c47cc108ca314dfd5cbb59cd3a143fd623", size = 1165527 },
+    { url = "https://files.pythonhosted.org/packages/41/05/d540049b1832d1062510efc6829634b7fbef5394c757d8312414fb65a3cb/torch-2.4.1-cp310-cp310-manylinux1_x86_64.whl", hash = "sha256:362f82e23a4cd46341daabb76fba08f04cd646df9bfaf5da50af97cb60ca4971", size = 797072810 },
+    { url = "https://files.pythonhosted.org/packages/a0/12/2162df9c47386ae7cedbc938f9703fee4792d93504fab8608d541e71ece3/torch-2.4.1-cp310-cp310-manylinux2014_aarch64.whl", hash = "sha256:e8ac1985c3ff0f60d85b991954cfc2cc25f79c84545aead422763148ed2759e3", size = 89699259 },
+    { url = "https://files.pythonhosted.org/packages/5d/4c/b2a59ff0e265f5ee154f0d81e948b1518b94f545357731e1a3245ee5d45b/torch-2.4.1-cp310-cp310-win_amd64.whl", hash = "sha256:91e326e2ccfb1496e3bee58f70ef605aeb27bd26be07ba64f37dcaac3d070ada", size = 199433813 },
+    { url = "https://files.pythonhosted.org/packages/dc/fb/1333ba666bbd53846638dd75a7a1d4eaf964aff1c482fc046e2311a1b499/torch-2.4.1-cp310-none-macosx_11_0_arm64.whl", hash = "sha256:d36a8ef100f5bff3e9c3cea934b9e0d7ea277cb8210c7152d34a9a6c5830eadd", size = 62139309 },
 ]
 
 [[package]]
@@ -940,13 +770,13 @@ wheels = [
 
 [[package]]
 name = "triton"
-version = "2.2.0"
+version = "3.0.0"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "filelock" },
 ]
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/95/05/ed974ce87fe8c8843855daa2136b3409ee1c126707ab54a8b72815c08b49/triton-2.2.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a2294514340cfe4e8f4f9e5c66c702744c4a117d25e618bd08469d0bfed1e2e5", size = 167900779 },
+    { url = "https://files.pythonhosted.org/packages/45/27/14cc3101409b9b4b9241d2ba7deaa93535a217a211c86c4cc7151fb12181/triton-3.0.0-1-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:e1efef76935b2febc365bfadf74bcb65a6f959a9872e5bddf44cc9e0adce1e1a", size = 209376304 },
 ]
 
 [[package]]
@@ -968,17 +798,12 @@ wheels = [
 ]
 
 [[package]]
-name = "xformers"
-version = "0.0.25.post1"
+name = "wheel"
+version = "0.44.0"
 source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "numpy" },
-    { name = "torch" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/ca/66/f5977922658ff2aea8b3222901fa0b9922778988c6d3a428cfd083892191/xformers-0.0.25.post1.tar.gz", hash = "sha256:397430bd0162fd5a75eb8bc50b0ba242200881e48fd6404a19376f853f8c0444", size = 4083274 }
+sdist = { url = "https://files.pythonhosted.org/packages/b7/a0/95e9e962c5fd9da11c1e28aa4c0d8210ab277b1ada951d2aee336b505813/wheel-0.44.0.tar.gz", hash = "sha256:a29c3f2817e95ab89aa4660681ad547c0e9547f20e75b0562fe7723c9a2a9d49", size = 100733 }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/5f/9b/f781a50d965717a2a2ea2c8d15e0f30deec8f9751a9874e850ba9ab0fadc/xformers-0.0.25.post1-cp310-cp310-manylinux2014_x86_64.whl", hash = "sha256:cdfe9560848fa5ba75fc04d3da8803658e35997adc6075ee6bbf6d67c1f0fa5e", size = 222517307 },
-    { url = "https://files.pythonhosted.org/packages/25/19/301789926809dc167ac104c29e26703369b3d47e6c1a2861db9efccbdc10/xformers-0.0.25.post1-cp310-cp310-win_amd64.whl", hash = "sha256:ddc22273f2ff06b886d9e86f17997e4f1f3074fdeb5d46bcdf50b704430df528", size = 208697268 },
+    { url = "https://files.pythonhosted.org/packages/1b/d1/9babe2ccaecff775992753d8686970b1e2755d21c8a63be73aba7a4e7d77/wheel-0.44.0-py3-none-any.whl", hash = "sha256:2376a90c98cc337d18623527a97c31797bd02bad0033d41547043a1cbfbe448f", size = 67059 },
 ]
 
 [[package]]