cleanup model types

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

modules/control/units/controlnet.py

@@ -181,7 +181,7 @@ def api_list_models(model_type: str = None):
         model_list += list(predefined_qwen)
     if model_type == 'hunyuandit' or model_type == 'all':
         model_list += list(predefined_hunyuandit)
-    if model_type == 'z_image':
+    if model_type == 'zimage':
         model_list += list(predefined_zimage)
     model_list += sorted(find_models())
     return model_list
@@ -207,7 +207,7 @@ def list_models(refresh=False):
         models = ['None'] + list(predefined_qwen) + sorted(find_models())
     elif modules.shared.sd_model_type == 'hunyuandit':
         models = ['None'] + list(predefined_hunyuandit) + sorted(find_models())
-    elif modules.shared.sd_model_type == 'z_image':
+    elif modules.shared.sd_model_type == 'zimage':
         models = ['None'] + list(predefined_zimage) + sorted(find_models())
     else:
         log.warning(f'Control {what} model list failed: unknown model type')
@@ -273,7 +273,7 @@ class ControlNet():
         elif shared.sd_model_type == 'hunyuandit':
             from diffusers import HunyuanDiT2DControlNetModel as cls
             config = 'Tencent-Hunyuan/HunyuanDiT-v1.2-ControlNet-Diffusers-Canny'
-        elif shared.sd_model_type == 'z_image':
+        elif shared.sd_model_type == 'zimage':
             from diffusers import ZImageControlNetModel as cls
             if '2.0' in model_id:
                 config = 'hlky/Z-Image-Turbo-Fun-Controlnet-Union-2.0'

modules/modeldata.py

@@ -22,7 +22,7 @@ def get_model_type(pipe):
         model_type = 'sd' # instaflow is compatible with sd
     elif "AnimateDiffPipeline" in name:
         model_type = 'sd' # animatediff is compatible with sd
-    elif "Kandinsky5" in name:
+    elif "Kandinsky5" in name and '2I' in name:
         model_type = 'kandinsky5'
     elif "Kandinsky3" in name:
         model_type = 'kandinsky3'
@@ -41,7 +41,7 @@ def get_model_type(pipe):
     elif "Flux" in name or "Flex1" in name or "Flex2" in name:
         model_type = 'f1'
     elif "ZImage" in name or "Z-Image" in name:
-        model_type = 'z_image'
+        model_type = 'zimage'
     elif "Lumina2" in name:
         model_type = 'lumina2'
     elif "Lumina" in name:
@@ -79,10 +79,22 @@ def get_model_type(pipe):
     elif 'LongCat' in name:
         model_type = 'longcat'
     elif 'GlmImage' in name:
-        model_type = 'glm_image'
+        model_type = 'glmimage'
     elif 'Ovis-Image' in name:
         model_type = 'ovis'
+    elif 'Wan' in name:
+        model_type = 'wanai'
+    elif 'ChronoEdit' in name:
+        model_type = 'chrono'
+    elif 'HDM-xut' in name:
+        model_type = 'hdm'
+    elif 'HunyuanImage3' in name:
+        model_type = 'hunyuanimage3'
+    elif 'HunyuanImage' in name:
+        model_type = 'hunyuanimage'
     # video models
+    elif "Kandinsky5" in name and '2V' in name:
+        model_type = 'kandinsky5video'
     elif "CogVideo" in name:
         model_type = 'cogvideo'
     elif 'HunyuanVideo15' in name:
@@ -95,17 +107,6 @@ def get_model_type(pipe):
         model_type = 'mochivideo'
     elif "Allegro" in name:
         model_type = 'allegrovideo'
-    # hybrid models
-    elif 'Wan' in name:
-        model_type = 'wanai'
-    elif 'ChronoEdit' in name:
-        model_type = 'chrono'
-    elif 'HDM-xut' in name:
-        model_type = 'hdm'
-    elif 'HunyuanImage3' in name:
-        model_type = 'hunyuanimage3'
-    elif 'HunyuanImage' in name:
-        model_type = 'hunyuanimage'
     # cloud models
     elif 'GoogleVeo' in name:
         model_type = 'veo3'

modules/sd_models.py

@@ -439,7 +439,7 @@ def load_diffuser_force(detected_model_type, checkpoint_info, diffusers_load_con
             from pipelines.model_kandinsky import load_kandinsky3
             sd_model = load_kandinsky3(checkpoint_info, diffusers_load_config)
             allow_post_quant = False
-        elif model_type in ['Kandinsky 5.0']:
+        elif model_type in ['Kandinsky 5.0'] and '2I' in model_type:
             from pipelines.model_kandinsky import load_kandinsky5
             sd_model = load_kandinsky5(checkpoint_info, diffusers_load_config)
             allow_post_quant = False

modules/sd_samplers_common.py

@@ -9,7 +9,7 @@ from modules import shared, devices, processing, images, sd_vae_approx, sd_vae_t
 
 SamplerData = namedtuple('SamplerData', ['name', 'constructor', 'aliases', 'options'])
 approximation_indexes = { "Simple": 0, "Approximate": 1, "TAESD": 2, "Full VAE": 3 }
-flow_models = ['f1', 'f2', 'sd3', 'lumina', 'auraflow', 'sana', 'z_image', 'lumina2', 'cogview4', 'h1', 'cosmos', 'chroma', 'omnigen', 'omnigen2', 'longcat']
+flow_models = ['f1', 'f2', 'sd3', 'lumina', 'auraflow', 'sana', 'zimage', 'lumina2', 'cogview4', 'h1', 'cosmos', 'chroma', 'omnigen', 'omnigen2', 'longcat']
 warned = False
 queue_lock = threading.Lock()
 

modules/sd_vae_remote.py

@@ -20,7 +20,7 @@ hf_decode_endpoints['auraflow'] = hf_decode_endpoints['sdxl']
 hf_decode_endpoints['omnigen'] = hf_decode_endpoints['sdxl']
 hf_decode_endpoints['h1'] = hf_decode_endpoints['f1']
 hf_decode_endpoints['chroma'] = hf_decode_endpoints['f1']
-hf_decode_endpoints['z_image'] = hf_decode_endpoints['f1']
+hf_decode_endpoints['zimage'] = hf_decode_endpoints['f1']
 hf_decode_endpoints['lumina2'] = hf_decode_endpoints['f1']
 
 hf_encode_endpoints = {
@@ -35,7 +35,7 @@ hf_encode_endpoints['hunyuandit'] = hf_encode_endpoints['sdxl']
 hf_encode_endpoints['auraflow'] = hf_encode_endpoints['sdxl']
 hf_encode_endpoints['omnigen'] = hf_encode_endpoints['sdxl']
 hf_encode_endpoints['h1'] = hf_encode_endpoints['f1']
-hf_encode_endpoints['z_image'] = hf_encode_endpoints['f1']
+hf_encode_endpoints['zimage'] = hf_encode_endpoints['f1']
 hf_encode_endpoints['lumina2'] = hf_encode_endpoints['f1']
 
 dtypes = {
@@ -92,7 +92,7 @@ def remote_decode(latents: torch.Tensor, width: int = 0, height: int = 0, model_
                 params["output_type"] = "pt"
                 params["output_tensor_type"] = "binary"
                 headers["Accept"] = "tensor/binary"
-            if model_type in {'f1', 'h1', 'z_image', 'lumina2', 'chroma'} and (width > 0) and (height > 0):
+            if model_type in {'f1', 'h1', 'zimage', 'lumina2', 'chroma'} and (width > 0) and (height > 0):
                 params['width'] = width
                 params['height'] = height
             if shared.sd_model.vae is not None and shared.sd_model.vae.config is not None:

modules/sd_vae_taesd.py

@@ -38,7 +38,7 @@ prev_cls = ''
 prev_type = ''
 prev_model = ''
 lock = threading.Lock()
-supported = ['sd', 'sdxl', 'sd3', 'f1', 'h1', 'z_image', 'lumina2', 'hunyuanvideo', 'wanai', 'chrono', 'mochivideo', 'pixartsigma', 'pixartalpha', 'hunyuandit', 'omnigen', 'qwen', 'longcat', 'glm_image']
+supported = ['sd', 'sdxl', 'sd3', 'f1', 'h1', 'zimage', 'lumina2', 'hunyuanvideo', 'wanai', 'chrono', 'mochivideo', 'pixartsigma', 'pixartalpha', 'hunyuandit', 'omnigen', 'qwen', 'longcat', 'glmimage']
 
 
 def warn_once(msg, variant=None):
@@ -59,7 +59,7 @@ def get_model(model_type = 'decoder', variant = None):
         model_cls = 'sd'
     elif model_cls in {'pixartsigma', 'hunyuandit', 'omnigen', 'auraflow'}:
         model_cls = 'sdxl'
-    elif model_cls in {'h1', 'z_image', 'lumina2', 'chroma', 'longcat', 'glm_image'}:
+    elif model_cls in {'h1', 'zimage', 'lumina2', 'chroma', 'longcat', 'glmimage'}:
         model_cls = 'f1'
     elif model_cls in {'wanai', 'qwen', 'chrono'}:
         variant = variant or 'TAE WanVideo'