From 67078318de898ede2baa5bea8d4017748aa85b31 Mon Sep 17 00:00:00 2001
From: Zhimin Li <46835311+zml-ai@users.noreply.github.com>
Date: Sat, 15 Jun 2024 01:29:43 +0800
Subject: [PATCH] fix trt bug

---
 hydit/inference.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/hydit/inference.py b/hydit/inference.py
index c7bcc15..682f29d 100644
--- a/hydit/inference.py
+++ b/hydit/inference.py
@@ -15,7 +15,7 @@
 from transformers import BertModel, BertTokenizer
 from transformers.modeling_utils import logger as tf_logger
 
-from .constants import SAMPLER_FACTORY, NEGATIVE_PROMPT
+from .constants import SAMPLER_FACTORY, NEGATIVE_PROMPT, TRT_MAX_WIDTH, TRT_MAX_HEIGHT, TRT_MAX_BATCH_SIZE
 from .diffusion.pipeline import StableDiffusionPipeline
 from .modules.models import HunYuanDiT, HUNYUAN_DIT_CONFIG
 from .modules.posemb_layers import get_2d_rotary_pos_embed, get_fill_resize_and_crop