draft schedule

Author: zRzRzRzRzRzRzR

scripts/convert_cogview4_to_diffusers.py

@@ -31,11 +31,10 @@
 from accelerate import init_empty_weights
 from transformers import PreTrainedTokenizerFast, GlmForCausalLM
 
-from diffusers import AutoencoderKL, CogVideoXDDIMScheduler, CogView4Pipeline, CogView3PlusTransformer2DModel
+from diffusers import AutoencoderKL, CogView4DDIMScheduler, CogView4Pipeline, CogView3PlusTransformer2DModel
 from diffusers.loaders.single_file_utils import convert_ldm_vae_checkpoint
 from diffusers.utils.import_utils import is_accelerate_available
 
-
 CTX = init_empty_weights if is_accelerate_available() else nullcontext
 
 parser = argparse.ArgumentParser()
@@ -170,16 +169,16 @@ def main(args):
             args.transformer_checkpoint_path
         )
         transformer = CogView3PlusTransformer2DModel(
-            patch_size = 2,
-            in_channels = 16,
-            num_layers = 28,
-            attention_head_dim= 128,
-            num_attention_heads = 32,
-            out_channels = 16,
-            text_embed_dim= 4096,
-            time_embed_dim = 512,
-            condition_dim= 256,
-            pos_embed_max_size = 128,
+            patch_size=2,
+            in_channels=16,
+            num_layers=28,
+            attention_head_dim=128,
+            num_attention_heads=32,
+            out_channels=16,
+            text_embed_dim=4096,
+            time_embed_dim=512,
+            condition_dim=256,
+            pos_embed_max_size=128,
         )
         transformer.load_state_dict(converted_transformer_state_dict, strict=True)
         if dtype is not None:
@@ -210,16 +209,20 @@ def main(args):
         if dtype is not None:
             vae = vae.to(dtype=dtype)
 
-    text_encoder_id = 'THUDM/glm-4-9b-hf'
+    text_encoder_id = "THUDM/glm-4-9b-hf"
     tokenizer = PreTrainedTokenizerFast.from_pretrained(text_encoder_id)
-    text_encoder = GlmForCausalLM.from_pretrained(text_encoder_id, cache_dir=args.text_encoder_cache_dir, torch_dtype=torch.bfloat16 if dtype=="bf16" else torch.float32)
+    text_encoder = GlmForCausalLM.from_pretrained(
+        text_encoder_id,
+        cache_dir=args.text_encoder_cache_dir,
+        torch_dtype=torch.bfloat16 if args.dtype == "bf16" else torch.float32,
+    )
     # Apparently, the conversion does not work anymore without this :shrug:
     for param in text_encoder.parameters():
         param.data = param.data.contiguous()
 
-    scheduler = CogVideoXDDIMScheduler.from_config(
+    scheduler = CogView4DDIMScheduler.from_config(
         {
-            "snr_shift_scale": 4.0,
+            "shift_scale": 1.0,
             "beta_end": 0.012,
             "beta_schedule": "scaled_linear",
             "beta_start": 0.00085,

src/diffusers/__init__.py

@@ -175,6 +175,7 @@
             "CMStochasticIterativeScheduler",
             "CogVideoXDDIMScheduler",
             "CogVideoXDPMScheduler",
+            "CogView4DDIMScheduler",
             "DDIMInverseScheduler",
             "DDIMParallelScheduler",
             "DDIMScheduler",
@@ -684,6 +685,7 @@
             CMStochasticIterativeScheduler,
             CogVideoXDDIMScheduler,
             CogVideoXDPMScheduler,
+            CogView4DDIMScheduler,
             DDIMInverseScheduler,
             DDIMParallelScheduler,
             DDIMScheduler,

src/diffusers/pipelines/cogview4/pipeline_cogview4.py

@@ -23,7 +23,7 @@
 from ...image_processor import VaeImageProcessor
 from ...models import AutoencoderKL, CogView3PlusTransformer2DModel
 from ...pipelines.pipeline_utils import DiffusionPipeline
-from ...schedulers import CogVideoXDDIMScheduler, CogVideoXDPMScheduler
+from ...schedulers import CogView4DDIMScheduler
 from ...utils import is_torch_xla_available, logging, replace_example_docstring
 from ...utils.torch_utils import randn_tensor
 from .pipeline_output import CogView4PipelineOutput
@@ -151,7 +151,7 @@ def __init__(
         text_encoder: GlmModel,
         vae: AutoencoderKL,
         transformer: CogView3PlusTransformer2DModel,
-        scheduler: Union[CogVideoXDDIMScheduler, CogVideoXDPMScheduler],
+        scheduler: CogView4DDIMScheduler,
     ):
         super().__init__()
 
@@ -318,7 +318,6 @@ def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype
             latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
         else:
             latents = latents.to(device)
-
         # scale the initial noise by the standard deviation required by the scheduler
         latents = latents * self.scheduler.init_noise_sigma
         return latents
@@ -517,8 +516,8 @@ def __call__(
         Examples:
 
         Returns:
-            [`~pipelines.cogview3.pipeline_CogView4.CogView3PipelineOutput`] or `tuple`:
-            [`~pipelines.cogview3.pipeline_CogView4.CogView3PipelineOutput`] if `return_dict` is True, otherwise a
+            [`~pipelines.cogview4.pipeline_CogView4.CogView3PipelineOutput`] or `tuple`:
+            [`~pipelines.cogview4.pipeline_CogView4.CogView3PipelineOutput`] if `return_dict` is True, otherwise a
             `tuple`. When returning a tuple, the first element is a list with the generated images.
         """
 
@@ -640,15 +639,13 @@ def __call__(
                     noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)
 
                 # compute the previous noisy sample x_t -> x_t-1
-                if not isinstance(self.scheduler, CogVideoXDPMScheduler):
+                if not isinstance(self.scheduler, CogView4DDIMScheduler):
                     latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
                 else:
                     latents, old_pred_original_sample = self.scheduler.step(
-                        noise_pred,
-                        old_pred_original_sample,
-                        t,
-                        timesteps[i - 1] if i > 0 else None,
-                        latents,
+                        model_output=noise_pred,
+                        timestep=t,
+                        sample=latents,
                         **extra_step_kwargs,
                         return_dict=False,
                     )

src/diffusers/schedulers/__init__.py

@@ -44,6 +44,7 @@
     _import_structure["scheduling_consistency_models"] = ["CMStochasticIterativeScheduler"]
     _import_structure["scheduling_ddim"] = ["DDIMScheduler"]
     _import_structure["scheduling_ddim_cogvideox"] = ["CogVideoXDDIMScheduler"]
+    _import_structure["scheduling_ddim_cogview4"] = ["CogView4DDIMScheduler"]
     _import_structure["scheduling_ddim_inverse"] = ["DDIMInverseScheduler"]
     _import_structure["scheduling_ddim_parallel"] = ["DDIMParallelScheduler"]
     _import_structure["scheduling_ddpm"] = ["DDPMScheduler"]
@@ -144,6 +145,7 @@
         from .scheduling_consistency_models import CMStochasticIterativeScheduler
         from .scheduling_ddim import DDIMScheduler
         from .scheduling_ddim_cogvideox import CogVideoXDDIMScheduler
+        from .scheduling_ddim_cogview4 import CogView4DDIMScheduler
         from .scheduling_ddim_inverse import DDIMInverseScheduler
         from .scheduling_ddim_parallel import DDIMParallelScheduler
         from .scheduling_ddpm import DDPMScheduler