From d6a8d2b1736fe94d1d0bdfdc7a2e4a536d4b2c03 Mon Sep 17 00:00:00 2001
From: Vladimir Mandic <mandic00@live.com>
Date: Thu, 26 Jun 2025 15:00:27 -0400
Subject: [PATCH] cleanup

Signed-off-by: Vladimir Mandic <mandic00@live.com>
---
 modules/processing_args.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/modules/processing_args.py b/modules/processing_args.py
index 42a8f6af0..bb09cc0a4 100644
--- a/modules/processing_args.py
+++ b/modules/processing_args.py
@@ -301,7 +301,7 @@ def set_pipeline_args(p, model, prompts:list, negative_prompts:list, prompts_2:t
             args['control_strength'] = p.denoising_strength
             args['width'] = p.width
             args['height'] = p.height
-    if 'Cosmos2TextToImagePipeline':
+    if 'Cosmos2' in model.__class__.__name__:
         kwargs['output_type'] = 'np' # cosmos uses wan-vae which is weird
     # set callbacks
     if 'prior_callback_steps' in possible:  # Wuerstchen / Cascade