Refactor and add comment

ROCm · Dec 17, 2024 · 9a9a849 · 9a9a849
1 parent 9d48096
commit 9a9a849
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/examples/diffusion/python_stable_diffusion_3/txt2img.py b/examples/diffusion/python_stable_diffusion_3/txt2img.py
@@ -485,10 +485,9 @@ def encode_token_weights(self, model_name, token_weight_pairs):
     def get_embeddings(self, prompt_tokens):
         l_out, l_pooled = self.encode_token_weights("clip-l",
                                                     prompt_tokens["l"])
+        # stable-diffusion-3-lite-onnx has swapped outputs for clip-l text encoder
         if l_out.shape != (1, 77, 768):
-            tmp = l_pooled
-            l_pooled = l_out
-            l_out = tmp
+            l_out, l_pooled = l_pooled, l_out
 
         g_out, g_pooled = self.encode_token_weights("clip-g",
                                                     prompt_tokens["g"])