fix trt bug

Tencent · Jun 14, 2024 · 6707831 · 6707831
1 parent 24c1c9a
commit 6707831
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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