From 816c7dc16c244a782351695946b810bf4dd90f0f Mon Sep 17 00:00:00 2001
From: chenyijian <chenyijian@chenyijiandeMacBook-Pro.local>
Date: Fri, 14 Jun 2024 11:55:41 +0800
Subject: [PATCH] fix bugs in trt.py line 353: current_unet is none not have
 cnets

---
 scripts/trt.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/scripts/trt.py b/scripts/trt.py
index ff0bbc7..44203f1 100644
--- a/scripts/trt.py
+++ b/scripts/trt.py
@@ -349,8 +349,10 @@ def process_batch(self, p, *args, **kwargs):
 
         if hasattr(p, "controlnet") and sd_unet.current_unet is not None:
             sd_unet.current_unet.cnets = p.controlnet
-        else:
+        elif sd_unet.current_unet is not None:
             sd_unet.current_unet.cnets = None
+        else:
+            sd_unet.current_unet = None
 
     def before_hr(self, p, *args):
         if self.idx != self.hr_idx: