[Pipiline] Wuerstchen v3 aka Stable Cascasde pipeline

scripts/convert_wuerstchen3.py

@@ -0,0 +1,176 @@
+# Run this script to convert the Wuerstchen V3 model weights to a diffusers pipeline.
+
+import torch
+from transformers import (
+    AutoTokenizer,
+    CLIPConfig,
+    CLIPImageProcessor,
+    CLIPTextModelWithProjection,
+    CLIPVisionModelWithProjection,
+)
+
+# from vqgan import VQModel
+from diffusers import (
+    DDPMWuerstchenScheduler,
+    WuerstchenV3DecoderPipeline,
+    WuerstchenV3PriorPipeline,
+)
+from diffusers.pipelines.wuerstchen import PaellaVQModel
+from diffusers.pipelines.wuerstchen3 import WuerstchenV3Unet
+
+
+device = "cpu"
+
+# set paths to model weights
+model_path = "../Wuerstchen"
+prior_checkpoint_path = f"{model_path}/v1.pt"
+decoder_checkpoint_path = f"{model_path}/base_120k.pt"
+
+
+# Clip Text encoder and tokenizer
+config = CLIPConfig.from_pretrained("laion/CLIP-ViT-bigG-14-laion2B-39B-b160k")
+config.text_config.projection_dim = config.projection_dim
+text_encoder = CLIPTextModelWithProjection.from_pretrained(
+    "laion/CLIP-ViT-bigG-14-laion2B-39B-b160k", config=config.text_config
+)
+tokenizer = AutoTokenizer.from_pretrained("laion/CLIP-ViT-bigG-14-laion2B-39B-b160k")
+
+# image processor
+feature_extractor = CLIPImageProcessor()
+image_encoder = CLIPVisionModelWithProjection.from_pretrained("openai/clip-vit-large-patch14")
+
+# Prior
+orig_state_dict = torch.load(prior_checkpoint_path, map_location=device)
+state_dict = {}
+for key in orig_state_dict.keys():
+    if key.endswith("in_proj_weight"):
+        weights = orig_state_dict[key].chunk(3, 0)
+        state_dict[key.replace("attn.in_proj_weight", "to_q.weight")] = weights[0]
+        state_dict[key.replace("attn.in_proj_weight", "to_k.weight")] = weights[1]
+        state_dict[key.replace("attn.in_proj_weight", "to_v.weight")] = weights[2]
+    elif key.endswith("in_proj_bias"):
+        weights = orig_state_dict[key].chunk(3, 0)
+        state_dict[key.replace("attn.in_proj_bias", "to_q.bias")] = weights[0]
+        state_dict[key.replace("attn.in_proj_bias", "to_k.bias")] = weights[1]
+        state_dict[key.replace("attn.in_proj_bias", "to_v.bias")] = weights[2]
+    elif key.endswith("out_proj.weight"):
+        weights = orig_state_dict[key]
+        state_dict[key.replace("attn.out_proj.weight", "to_out.0.weight")] = weights
+    elif key.endswith("out_proj.bias"):
+        weights = orig_state_dict[key]
+        state_dict[key.replace("attn.out_proj.bias", "to_out.0.bias")] = weights
+    else:
+        state_dict[key] = orig_state_dict[key]
+prior_model = WuerstchenV3Unet(
+    c_in=16,
+    c_out=16,
+    c_r=64,
+    patch_size=1,
+    c_cond=2048,
+    c_hidden=[2048, 2048],
+    nhead=[32, 32],
+    blocks=[[8, 24], [24, 8]],
+    block_repeat=[[1, 1], [1, 1]],
+    level_config=["CTA", "CTA"],
+    c_clip_text=1280,
+    c_clip_text_pooled=1280,
+    c_clip_img=768,
+    c_clip_seq=4,
+    kernel_size=3,
+    dropout=[0.1, 0.1],
+    self_attn=True,
+    t_conds=["sca", "crp"],
+    switch_level=[False],
+).to(device)
+prior_model.load_state_dict(state_dict)
+
+# scheduler for prior and decoder
+scheduler = DDPMWuerstchenScheduler()
+
+# Prior pipeline
+prior_pipeline = WuerstchenV3PriorPipeline(
+    prior=prior_model,
+    tokenizer=tokenizer,
+    text_encoder=text_encoder,
+    image_encoder=image_encoder,
+    scheduler=scheduler,
+    feature_extractor=feature_extractor,
+)
+prior_pipeline.save_pretrained("wuerstchenV3-prior")
+
+# Decoder
+orig_state_dict = torch.load(decoder_checkpoint_path, map_location=device)
+state_dict = {}
+for key in orig_state_dict.keys():
+    if key.endswith("in_proj_weight"):
+        weights = orig_state_dict[key].chunk(3, 0)
+        state_dict[key.replace("attn.in_proj_weight", "to_q.weight")] = weights[0]
+        state_dict[key.replace("attn.in_proj_weight", "to_k.weight")] = weights[1]
+        state_dict[key.replace("attn.in_proj_weight", "to_v.weight")] = weights[2]
+    elif key.endswith("in_proj_bias"):
+        weights = orig_state_dict[key].chunk(3, 0)
+        state_dict[key.replace("attn.in_proj_bias", "to_q.bias")] = weights[0]
+        state_dict[key.replace("attn.in_proj_bias", "to_k.bias")] = weights[1]
+        state_dict[key.replace("attn.in_proj_bias", "to_v.bias")] = weights[2]
+    elif key.endswith("out_proj.weight"):
+        weights = orig_state_dict[key]
+        state_dict[key.replace("attn.out_proj.weight", "to_out.0.weight")] = weights
+    elif key.endswith("out_proj.bias"):
+        weights = orig_state_dict[key]
+        state_dict[key.replace("attn.out_proj.bias", "to_out.0.bias")] = weights
+    # rename clip_mapper to clip_txt_pooled_mapper
+    elif key.endswith("clip_mapper.weight"):
+        weights = orig_state_dict[key]
+        state_dict[key.replace("clip_mapper.weight", "clip_txt_pooled_mapper.weight")] = weights
+    elif key.endswith("clip_mapper.bias"):
+        weights = orig_state_dict[key]
+        state_dict[key.replace("clip_mapper.bias", "clip_txt_pooled_mapper.bias")] = weights
+    else:
+        state_dict[key] = orig_state_dict[key]
+decoder = WuerstchenV3Unet(
+    c_in=4,
+    c_out=4,
+    c_r=64,
+    patch_size=2,
+    c_cond=1280,
+    c_hidden=[320, 640, 1280, 1280],
+    nhead=[-1, -1, 20, 20],
+    blocks=[[2, 6, 28, 6], [6, 28, 6, 2]],
+    block_repeat=[[1, 1, 1, 1], [3, 3, 2, 2]],
+    level_config=["CT", "CT", "CTA", "CTA"],
+    c_clip_text_pooled=1280,
+    c_clip_seq=4,
+    c_effnet=16,
+    c_pixels=3,
+    kernel_size=3,
+    dropout=[0, 0, 0.1, 0.1],
+    self_attn=True,
+    t_conds=["sca"],
+).to(device)
+decoder.load_state_dict(state_dict)
+
+# VQGAN from V2
+vqmodel = PaellaVQModel.from_pretrained("warp-ai/wuerstchen", subfolder="vqgan")
+
+# Decoder pipeline
+decoder_pipeline = WuerstchenV3DecoderPipeline(
+    decoder=decoder, text_encoder=text_encoder, tokenizer=tokenizer, vqgan=vqmodel, scheduler=scheduler
+)
+decoder_pipeline.save_pretrained("wuerstchenV3")
+
+# TODO
+# # Wuerstchen pipeline
+# wuerstchen_pipeline = WuerstchenCombinedPipeline(
+#     # Decoder
+#     text_encoder=gen_text_encoder,
+#     tokenizer=gen_tokenizer,
+#     decoder=decoder,
+#     scheduler=scheduler,
+#     vqgan=vqmodel,
+#     # Prior
+#     prior_tokenizer=tokenizer,
+#     prior_text_encoder=text_encoder,
+#     prior=prior_model,
+#     prior_scheduler=scheduler,
+# )
+# wuerstchen_pipeline.save_pretrained("warp-ai/WuerstchenCombinedPipeline")

src/diffusers/__init__.py

@@ -307,6 +307,9 @@
             "WuerstchenCombinedPipeline",
             "WuerstchenDecoderPipeline",
             "WuerstchenPriorPipeline",
+            "WuerstchenV3CombinedPipeline",
+            "WuerstchenV3DecoderPipeline",
+            "WuerstchenV3PriorPipeline",
         ]
     )
 
@@ -670,6 +673,9 @@
             WuerstchenCombinedPipeline,
             WuerstchenDecoderPipeline,
             WuerstchenPriorPipeline,
+            WuerstchenV3CombinedPipeline,
+            WuerstchenV3DecoderPipeline,
+            WuerstchenV3PriorPipeline,
         )
 
     try:

src/diffusers/pipelines/__init__.py

@@ -233,6 +233,11 @@
         "WuerstchenDecoderPipeline",
         "WuerstchenPriorPipeline",
     ]
+    _import_structure["wuerstchen3"] = [
+        "WuerstchenV3CombinedPipeline",
+        "WuerstchenV3DecoderPipeline",
+        "WuerstchenV3PriorPipeline",
+    ]
 try:
     if not is_onnx_available():
         raise OptionalDependencyNotAvailable()
@@ -471,6 +476,11 @@
             WuerstchenDecoderPipeline,
             WuerstchenPriorPipeline,
         )
+        from .wuerstchen3 import (
+            WuerstchenV3CombinedPipeline,
+            WuerstchenV3DecoderPipeline,
+            WuerstchenV3PriorPipeline,
+        )
 
         try:
             if not is_onnx_available():

src/diffusers/pipelines/wuerstchen/modeling_wuerstchen_common.py

@@ -32,13 +32,20 @@ def forward(self, x):
 
 
 class TimestepBlock(nn.Module):
-    def __init__(self, c, c_timestep):
+    def __init__(self, c, c_timestep, conds=[]):
         super().__init__()
         linear_cls = nn.Linear if USE_PEFT_BACKEND else LoRACompatibleLinear
         self.mapper = linear_cls(c_timestep, c * 2)
+        self.conds = conds
+        for cname in conds:
+            setattr(self, f"mapper_{cname}", linear_cls(c_timestep, c * 2))
 
     def forward(self, x, t):
-        a, b = self.mapper(t)[:, :, None, None].chunk(2, dim=1)
+        t = t.chunk(len(self.conds) + 1, dim=1)
+        a, b = self.mapper(t[0])[:, :, None, None].chunk(2, dim=1)
+        for i, c in enumerate(self.conds):
+            ac, bc = getattr(self, f"mapper_{c}")(t[i + 1])[:, :, None, None].chunk(2, dim=1)
+            a, b = a + ac, b + bc
         return x * (1 + a) + b
 
 
@@ -49,10 +56,14 @@ def __init__(self, c, c_skip=0, kernel_size=3, dropout=0.0):
         conv_cls = nn.Conv2d if USE_PEFT_BACKEND else LoRACompatibleConv
         linear_cls = nn.Linear if USE_PEFT_BACKEND else LoRACompatibleLinear
 
-        self.depthwise = conv_cls(c + c_skip, c, kernel_size=kernel_size, padding=kernel_size // 2, groups=c)
+        self.depthwise = conv_cls(c, c, kernel_size=kernel_size, padding=kernel_size // 2, groups=c)
         self.norm = WuerstchenLayerNorm(c, elementwise_affine=False, eps=1e-6)
         self.channelwise = nn.Sequential(
-            linear_cls(c, c * 4), nn.GELU(), GlobalResponseNorm(c * 4), nn.Dropout(dropout), linear_cls(c * 4, c)
+            linear_cls(c + c_skip, c * 4),
+            nn.GELU(),
+            GlobalResponseNorm(c * 4),
+            nn.Dropout(dropout),
+            linear_cls(c * 4, c),
         )
 
     def forward(self, x, x_skip=None):

src/diffusers/pipelines/wuerstchen/modeling_wuerstchen_diffnext.py

@@ -233,7 +233,7 @@ def forward(self, x, r, effnet, clip=None, x_cat=None, eps=1e-3, return_noise=Tr
 
 
 class ResBlockStageB(nn.Module):
-    def __init__(self, c, c_skip=None, kernel_size=3, dropout=0.0):
+    def __init__(self, c, c_skip=0, kernel_size=3, dropout=0.0):
         super().__init__()
         self.depthwise = nn.Conv2d(c, c, kernel_size=kernel_size, padding=kernel_size // 2, groups=c)
         self.norm = WuerstchenLayerNorm(c, elementwise_affine=False, eps=1e-6)

src/diffusers/pipelines/wuerstchen/pipeline_wuerstchen.py

@@ -423,9 +423,9 @@ def __call__(
             latents = self.vqgan.config.scale_factor * latents
             images = self.vqgan.decode(latents).sample.clamp(0, 1)
             if output_type == "np":
-                images = images.permute(0, 2, 3, 1).cpu().numpy()
+                images = images.permute(0, 2, 3, 1).float().cpu().numpy()
             elif output_type == "pil":
-                images = images.permute(0, 2, 3, 1).cpu().numpy()
+                images = images.permute(0, 2, 3, 1).float().cpu().numpy()
                 images = self.numpy_to_pil(images)
         else:
             images = latents

src/diffusers/pipelines/wuerstchen3/__init__.py

@@ -0,0 +1,52 @@
+from typing import TYPE_CHECKING
+
+from ...utils import (
+    DIFFUSERS_SLOW_IMPORT,
+    OptionalDependencyNotAvailable,
+    _LazyModule,
+    get_objects_from_module,
+    is_torch_available,
+    is_transformers_available,
+)
+
+
+_dummy_objects = {}
+_import_structure = {}
+
+try:
+    if not (is_transformers_available() and is_torch_available()):
+        raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+    from ...utils import dummy_torch_and_transformers_objects
+
+    _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects))
+else:
+    _import_structure["modeling_wuerstchen3_common"] = ["WuerstchenV3Unet"]
+    _import_structure["pipeline_wuerstchen3"] = ["WuerstchenV3DecoderPipeline"]
+    _import_structure["pipeline_wuerstchen3_combined"] = ["WuerstchenV3CombinedPipeline"]
+    _import_structure["pipeline_wuerstchen3_prior"] = ["WuerstchenV3PriorPipeline"]
+
+
+if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
+    try:
+        if not (is_transformers_available() and is_torch_available()):
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        from ...utils.dummy_torch_and_transformers_objects import *  # noqa F403
+    else:
+        from .modeling_wuerstchen3_common import WuerstchenV3Unet
+        from .pipeline_wuerstchen3 import WuerstchenV3DecoderPipeline
+        from .pipeline_wuerstchen3_combined import WuerstchenV3CombinedPipeline
+        from .pipeline_wuerstchen3_prior import WuerstchenV3PriorPipeline
+else:
+    import sys
+
+    sys.modules[__name__] = _LazyModule(
+        __name__,
+        globals()["__file__"],
+        _import_structure,
+        module_spec=__spec__,
+    )
+
+    for name, value in _dummy_objects.items():
+        setattr(sys.modules[__name__], name, value)