test

lanluo-nvidia · lanluo-nvidia · commit 2f4afae992d9 · 2025-07-03T17:15:07.000-07:00
diff --git a/examples/apps/flux_demo.py b/examples/apps/flux_demo.py
@@ -10,6 +10,7 @@
 import torch_tensorrt
 from accelerate.hooks import remove_hook_from_module
 from diffusers import FluxPipeline
+from diffusers.models.transformers.transformer_flux import FluxTransformer2DModel
 
 DEVICE = "cuda:0"
 
