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