add cogview4

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

modules/model_cogview.py

@@ -0,0 +1,95 @@
+import transformers
+import diffusers
+from modules import shared, devices, sd_models
+
+
+def load_common(diffusers_load_config={}, module=None):
+    from modules import model_quant, modelloader
+    modelloader.hf_login()
+
+    if 'torch_dtype' not in diffusers_load_config:
+        diffusers_load_config['torch_dtype'] = 'torch.float16'
+    if 'low_cpu_mem_usage' in diffusers_load_config:
+        del diffusers_load_config['low_cpu_mem_usage']
+    if 'load_connected_pipeline' in diffusers_load_config:
+        del diffusers_load_config['load_connected_pipeline']
+    if 'safety_checker' in diffusers_load_config:
+        del diffusers_load_config['safety_checker']
+    if 'requires_safety_checker' in diffusers_load_config:
+        del diffusers_load_config['requires_safety_checker']
+
+    quant_args = {}
+    if not quant_args:
+        quant_args = model_quant.create_bnb_config(quant_args, module=module)
+    if not quant_args:
+        quant_args = model_quant.create_ao_config(quant_args, module=module)
+    if quant_args:
+        shared.log.debug(f'Load model: type=CogView quantization module="{module}" {quant_args}')
+
+    return diffusers_load_config, quant_args
+
+
+def load_cogview3(checkpoint_info, diffusers_load_config={}):
+    repo_id = sd_models.path_to_repo(checkpoint_info.name)
+    shared.log.debug(f'Load model: type=CogView3 model="{checkpoint_info.name}" repo="{repo_id}" offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype}')
+
+    diffusers_load_config, quant_args = load_common(diffusers_load_config, module='Model')
+    transformer = diffusers.CogView3PlusTransformer2DModel.from_pretrained(
+        repo_id,
+        subfolder="transformer",
+        cache_dir=shared.opts.diffusers_dir,
+        **diffusers_load_config,
+        **quant_args,
+    )
+
+    diffusers_load_config, quant_args = load_common(diffusers_load_config, module='Text Encoder')
+    text_encoder = transformers.T5EncoderModel.from_pretrained(
+        repo_id,
+        subfolder="text_encoder",
+        cache_dir=shared.opts.diffusers_dir,
+        **diffusers_load_config,
+        **quant_args,
+    )
+
+    pipe = diffusers.CogView3PlusPipeline.from_pretrained(
+        repo_id,
+        text_encoder=text_encoder,
+        transformer=transformer,
+        cache_dir=shared.opts.diffusers_dir,
+        **diffusers_load_config,
+    )
+    devices.torch_gc()
+    return pipe
+
+
+def load_cogview4(checkpoint_info, diffusers_load_config={}):
+    repo_id = sd_models.path_to_repo(checkpoint_info.name)
+    shared.log.debug(f'Load model: type=CogView4 model="{checkpoint_info.name}" repo="{repo_id}" offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype}')
+
+    diffusers_load_config, quant_args = load_common(diffusers_load_config, module='Model')
+    transformer = diffusers.CogView4Transformer2DModel.from_pretrained(
+        repo_id,
+        subfolder="transformer",
+        cache_dir=shared.opts.diffusers_dir,
+        **diffusers_load_config,
+        **quant_args,
+    )
+
+    diffusers_load_config, quant_args = load_common(diffusers_load_config, module='Text Encoder')
+    text_encoder = transformers.T5EncoderModel.from_pretrained(
+        repo_id,
+        subfolder="text_encoder",
+        cache_dir=shared.opts.diffusers_dir,
+        **diffusers_load_config,
+        **quant_args,
+    )
+
+    pipe = diffusers.CogView4Pipeline.from_pretrained(
+        repo_id,
+        text_encoder=text_encoder,
+        transformer=transformer,
+        cache_dir=shared.opts.diffusers_dir,
+        **diffusers_load_config,
+    )
+    devices.torch_gc()
+    return pipe

modules/model_flux.py

@@ -110,39 +110,6 @@ def load_flux_bnb(checkpoint_info, diffusers_load_config): # pylint: disable=unu
     return transformer, text_encoder_2
 
 
-"""
-def quant_flux_bnb(checkpoint_info, transformer, text_encoder_2):
-    repo_id = sd_models.path_to_repo(checkpoint_info.name)
-    cache_dir=shared.opts.diffusers_dir
-    if len(shared.opts.bnb_quantization) > 0 and (transformer is None or text_encoder_2 is None):
-        from modules.model_quant import load_bnb
-        load_bnb('Load model: type=FLUX')
-        try:
-            bnb_config = diffusers.BitsAndBytesConfig(
-                load_in_8bit=shared.opts.bnb_quantization_type in ['fp8'],
-                load_in_4bit=shared.opts.bnb_quantization_type in ['nf4', 'fp4'],
-                bnb_4bit_quant_storage=shared.opts.bnb_quantization_storage,
-                bnb_4bit_quant_type=shared.opts.bnb_quantization_type,
-                bnb_4bit_compute_dtype=devices.dtype
-            )
-            if ('Model' in shared.opts.bnb_quantization) and (transformer is None):
-                transformer = diffusers.FluxTransformer2DModel.from_pretrained(repo_id, subfolder="transformer", cache_dir=cache_dir, quantization_config=bnb_config, torch_dtype=devices.dtype)
-                shared.log.debug(f'Quantization: module=transformer type=bnb dtype={shared.opts.bnb_quantization_type} storage={shared.opts.bnb_quantization_storage}')
-            if ('Text Encoder' in shared.opts.bnb_quantization) and (text_encoder_2 is None):
-                if repo_id == 'sayakpaul/flux.1-dev-nf4':
-                    repo_id = 'black-forest-labs/FLUX.1-dev' # workaround since sayakpaul model is missing model_index.json
-                text_encoder_2 = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder="text_encoder_2", cache_dir=cache_dir, quantization_config=bnb_config, torch_dtype=devices.dtype)
-                shared.log.debug(f'Quantization: module=t5 type=bnb dtype={shared.opts.bnb_quantization_type} storage={shared.opts.bnb_quantization_storage}')
-        except Exception as e:
-            shared.log.error(f"Load model: type=FLUX failed quantize using BnB: {e}")
-            transformer, text_encoder_2 = None, None
-            if debug:
-                from modules import errors
-                errors.display(e, 'FLUX:')
-    return transformer, text_encoder_2
-"""
-
-
 def load_quants(kwargs, repo_id, cache_dir, allow_quant):
     try:
         if not allow_quant:

modules/model_quant.py

@@ -30,10 +30,10 @@ def get_quant(name):
     return 'none'
 
 
-def create_bnb_config(kwargs = None, allow_bnb: bool = True):
+def create_bnb_config(kwargs = None, allow_bnb: bool = True, module: str = 'Model'):
     from modules import shared, devices
     if len(shared.opts.bnb_quantization) > 0 and allow_bnb:
-        if 'Model' in shared.opts.bnb_quantization:
+        if 'Model' in shared.opts.bnb_quantization or (module is not None and module in shared.opts.bnb_quantization):
             load_bnb()
             if bnb is None:
                 return kwargs
@@ -53,10 +53,10 @@ def create_bnb_config(kwargs = None, allow_bnb: bool = True):
     return kwargs
 
 
-def create_ao_config(kwargs = None, allow_ao: bool = True):
+def create_ao_config(kwargs = None, allow_ao: bool = True, module: str = 'Model'):
     from modules import shared
     if len(shared.opts.torchao_quantization) > 0 and shared.opts.torchao_quantization_mode == 'pre' and allow_ao:
-        if 'Model' in shared.opts.torchao_quantization:
+        if 'Model' in shared.opts.torchao_quantization or (module is not None and module in shared.opts.torchao_quantization):
             load_torchao()
             if ao is None:
                 return kwargs

modules/modeldata.py

@@ -37,6 +37,10 @@ def get_model_type(pipe):
         model_type = 'lumina'
     elif "OmniGen" in name:
         model_type = 'omnigen'
+    elif "CogView3" in name:
+        model_type = 'cogview3'
+    elif "CogView4" in name:
+        model_type = 'cogview4'
     elif "CogVideo" in name:
         model_type = 'cogvideox'
     elif "Sana" in name:

modules/sd_detect.py

@@ -75,8 +75,10 @@ def detect_pipeline(f: str, op: str = 'model', warning=True, quiet=False):
                 guess = 'Kolors'
             if 'auraflow' in f.lower():
                 guess = 'AuraFlow'
-            if 'cogview' in f.lower():
-                guess = 'CogView'
+            if 'cogview3' in f.lower():
+                guess = 'CogView3'
+            if 'cogview4' in f.lower():
+                guess = 'CogView4'
             if 'meissonic' in f.lower():
                 guess = 'Meissonic'
                 pipeline = 'custom'

modules/sd_models.py

@@ -295,9 +295,13 @@ def load_diffuser_force(model_type, checkpoint_info, diffusers_load_config, op='
             sd_model = load_lumina2(checkpoint_info, diffusers_load_config)
         elif model_type in ['Stable Diffusion 3']:
             from modules.model_sd3 import load_sd3
-            shared.log.debug(f'Load {op}: model="Stable Diffusion 3"')
-            shared.opts.scheduler = 'Default'
             sd_model = load_sd3(checkpoint_info, cache_dir=shared.opts.diffusers_dir, config=diffusers_load_config.get('config', None))
+        elif model_type in ['CogView3']: # forced pipeline
+            from modules.model_cogview import load_cogview3
+            sd_model = load_cogview3(checkpoint_info, diffusers_load_config)
+        elif model_type in ['CogView4']: # forced pipeline
+            from modules.model_cogview import load_cogview4
+            sd_model = load_cogview4(checkpoint_info, diffusers_load_config)
         elif model_type in ['Meissonic']: # forced pipeline
             from modules.model_meissonic import load_meissonic
             sd_model = load_meissonic(checkpoint_info, diffusers_load_config)

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', 'sd3', 'lumina', 'auraflow', 'sana', 'lumina2']
+flow_models = ['f1', 'sd3', 'lumina', 'auraflow', 'sana', 'lumina2', 'cogview4']
 warned = False
 queue_lock = threading.Lock()
 

modules/shared_items.py

@@ -89,7 +89,8 @@ def get_pipelines():
         'SegMoE': getattr(diffusers, 'StableDiffusionPipeline', None), # dynamically redefined and loaded in sd_models.load_diffuser
         'Kolors': getattr(diffusers, 'KolorsPipeline', None),
         'AuraFlow': getattr(diffusers, 'AuraFlowPipeline', None),
-        'CogView': getattr(diffusers, 'CogView3PlusPipeline', None),
+        'CogView3': getattr(diffusers, 'CogView3PlusPipeline', None),
+        'CogView4': getattr(diffusers, 'CogView4Pipeline', None),
         'Stable Cascade': getattr(diffusers, 'StableCascadeCombinedPipeline', None),
         'PixArt-Sigma': getattr(diffusers, 'PixArtSigmaPipeline', None),
         'HunyuanDiT': getattr(diffusers, 'HunyuanDiTPipeline', None),