cleanup

Signed-off-by: Vladimir Mandic <mandic00@live.com>

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