Merge remote-tracking branch 'origin/main'

conversion.py

-# coding=utf-8
-# Copyright 2022 HuggingFace Inc.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-
-import inspect
-import tempfile
-import unittest
-
-import numpy as np
-import torch
-
-from diffusers import (
-    AutoencoderKL,
-    DDIMPipeline,
-    DDIMScheduler,
-    DDPMPipeline,
-    DDPMScheduler,
-    GlidePipeline,
-    GlideSuperResUNetModel,
-    GlideTextToImageUNetModel,
-    LatentDiffusionPipeline,
-    LatentDiffusionUncondPipeline,
-    NCSNpp,
-    PNDMPipeline,
-    PNDMScheduler,
-    ScoreSdeVePipeline,
-    ScoreSdeVeScheduler,
-    ScoreSdeVpPipeline,
-    ScoreSdeVpScheduler,
-    UNetLDMModel,
-    UNetModel,
-    UNetUnconditionalModel,
-    VQModel,
-)
-from diffusers.configuration_utils import ConfigMixin
-from diffusers.pipeline_utils import DiffusionPipeline
-from diffusers.testing_utils import floats_tensor, slow, torch_device
-from diffusers.training_utils import EMAModel
-
-
-# 1. LDM
-
-def test_output_pretrained_ldm_dummy():
-    model = UNetUnconditionalModel.from_pretrained("fusing/unet-ldm-dummy", ldm=True)
-    model.eval()
-
-    torch.manual_seed(0)
-    if torch.cuda.is_available():
-        torch.cuda.manual_seed_all(0)
-
-    noise = torch.randn(1, model.config.in_channels, model.config.image_size, model.config.image_size)
-    time_step = torch.tensor([10] * noise.shape[0])
-
-    with torch.no_grad():
-        output = model(noise, time_step)
-        
-    print(model)
-    import ipdb; ipdb.set_trace()
-
-
-def test_output_pretrained_ldm():
-    model = UNetUnconditionalModel.from_pretrained("fusing/latent-diffusion-celeba-256", subfolder="unet", ldm=True)
-    model.eval()
-
-    torch.manual_seed(0)
-    if torch.cuda.is_available():
-        torch.cuda.manual_seed_all(0)
-
-    noise = torch.randn(1, model.config.in_channels, model.config.image_size, model.config.image_size)
-    time_step = torch.tensor([10] * noise.shape[0])
-
-    with torch.no_grad():
-        output = model(noise, time_step)
-            
-    print(model)
-    import ipdb; ipdb.set_trace()
-
-# To see the how the final model should look like
-# => this is the architecture in which the model should be saved in the new format
-# -> verify new repo with the following tests (in `test_modeling_utils.py`)
-# - test_ldm_uncond (in PipelineTesterMixin)
-# - test_output_pretrained  ( in UNetLDMModelTests)
-
-#test_output_pretrained_ldm_dummy()
-#test_output_pretrained_ldm()
-
-
-# 2. DDPM
-
-def get_model(model_id):
-    model = UNetUnconditionalModel.from_pretrained(model_id, ldm=True)
-
-    noise = torch.randn(1, model.config.in_channels, model.config.image_size, model.config.image_size)
-    time_step = torch.tensor([10] * noise.shape[0])
-
-    with torch.no_grad():
-        output = model(noise, time_step)
-
-    print(model)
-
-# Repos to convert and port to google (part of https://github.com/hojonathanho/diffusion)
-# - fusing/ddpm_dummy
-# - fusing/ddpm-cifar10
-# - https://huggingface.co/fusing/ddpm-lsun-church-ema
-# - https://huggingface.co/fusing/ddpm-lsun-bedroom-ema
-# - https://huggingface.co/fusing/ddpm-celeba-hq
-
-# tests to make sure to pass
-# - test_ddim_cifar10, test_ddim_lsun, test_ddpm_cifar10, test_ddim_cifar10 (in PipelineTesterMixin)
-# - test_output_pretrained  ( in UNetModelTests)
-
-# e.g.
-get_model("fusing/ddpm-cifar10")
-
-# 3. NCSNpp
-
-# Repos to convert and port to google (part of https://github.com/yang-song/score_sde)
-# - https://huggingface.co/fusing/ffhq_ncsnpp
-# - https://huggingface.co/fusing/church_256-ncsnpp-ve
-# - https://huggingface.co/fusing/celebahq_256-ncsnpp-ve
-# - https://huggingface.co/fusing/bedroom_256-ncsnpp-ve
-# - https://huggingface.co/fusing/ffhq_256-ncsnpp-ve
-
-# tests to make sure to pass
-# - test_score_sde_ve_pipeline (in PipelineTesterMixin)
-# - test_output_pretrained_ve_mid, test_output_pretrained_ve_large (in NCSNppModelTests)

scripts/convert_ldm_original_checkpoint_to_diffusers.py

@@ -17,6 +17,7 @@
 import argparse
 import json
 import torch
+from diffusers import VQModel, DDPMScheduler, UNetUnconditionalModel, LatentDiffusionUncondPipeline
 
 
 def shave_segments(path, n_shave_prefix_segments=1):
@@ -314,4 +315,18 @@ if __name__ == "__main__":
         config = json.loads(f.read())
 
     converted_checkpoint = convert_ldm_checkpoint(checkpoint, config)
-    torch.save(checkpoint, args.dump_path)
+
+    if "ldm" in config:
+        del config["ldm"]
+
+    model = UNetUnconditionalModel(**config)
+    model.load_state_dict(converted_checkpoint)
+
+    try:
+        scheduler = DDPMScheduler.from_config("/".join(args.checkpoint_path.split("/")[:-1]))
+        vqvae = VQModel.from_pretrained("/".join(args.checkpoint_path.split("/")[:-1]))
+
+        pipe = LatentDiffusionUncondPipeline(unet=model, scheduler=scheduler, vae=vqvae)
+        pipe.save_pretrained(args.dump_path)
+    except:
+        model.save_pretrained(args.dump_path)

scripts/convert_ncsnpp_original_checkpoint_to_diffusers.py

@@ -20,22 +20,21 @@ import torch
 from diffusers import UNetUnconditionalModel
 
 
-
 def convert_ncsnpp_checkpoint(checkpoint, config):
     """
     Takes a state dict and the path to
     """
     new_model_architecture = UNetUnconditionalModel(**config)
-    new_model_architecture.time_steps.W.data= checkpoint['all_modules.0.W'].data
-    new_model_architecture.time_steps.weight.data = checkpoint['all_modules.0.W'].data
-    new_model_architecture.time_embedding.linear_1.weight.data = checkpoint['all_modules.1.weight'].data
-    new_model_architecture.time_embedding.linear_1.bias.data = checkpoint['all_modules.1.bias'].data
+    new_model_architecture.time_steps.W.data = checkpoint["all_modules.0.W"].data
+    new_model_architecture.time_steps.weight.data = checkpoint["all_modules.0.W"].data
+    new_model_architecture.time_embedding.linear_1.weight.data = checkpoint["all_modules.1.weight"].data
+    new_model_architecture.time_embedding.linear_1.bias.data = checkpoint["all_modules.1.bias"].data
 
-    new_model_architecture.time_embedding.linear_2.weight.data = checkpoint['all_modules.2.weight'].data
-    new_model_architecture.time_embedding.linear_2.bias.data= checkpoint['all_modules.2.bias'].data
+    new_model_architecture.time_embedding.linear_2.weight.data = checkpoint["all_modules.2.weight"].data
+    new_model_architecture.time_embedding.linear_2.bias.data = checkpoint["all_modules.2.bias"].data
 
-    new_model_architecture.conv_in.weight.data = checkpoint['all_modules.3.weight'].data
-    new_model_architecture.conv_in.bias.data = checkpoint['all_modules.3.bias'].data
+    new_model_architecture.conv_in.weight.data = checkpoint["all_modules.3.weight"].data
+    new_model_architecture.conv_in.bias.data = checkpoint["all_modules.3.bias"].data
 
     new_model_architecture.conv_norm_out.weight.data = checkpoint[list(checkpoint.keys())[-4]].data
     new_model_architecture.conv_norm_out.bias.data = checkpoint[list(checkpoint.keys())[-3]].data
@@ -44,8 +43,7 @@ def convert_ncsnpp_checkpoint(checkpoint, config):
 
     module_index = 4
 
-
-    def set_attention_weights(new_layer,old_checkpoint,index):
+    def set_attention_weights(new_layer, old_checkpoint, index):
         new_layer.query.weight.data = old_checkpoint[f"all_modules.{index}.NIN_0.W"].data.T
         new_layer.key.weight.data = old_checkpoint[f"all_modules.{index}.NIN_1.W"].data.T
         new_layer.value.weight.data = old_checkpoint[f"all_modules.{index}.NIN_2.W"].data.T
@@ -60,7 +58,7 @@ def convert_ncsnpp_checkpoint(checkpoint, config):
         new_layer.group_norm.weight.data = old_checkpoint[f"all_modules.{index}.GroupNorm_0.weight"].data
         new_layer.group_norm.bias.data = old_checkpoint[f"all_modules.{index}.GroupNorm_0.bias"].data
 
-    def set_resnet_weights(new_layer,old_checkpoint,index):
+    def set_resnet_weights(new_layer, old_checkpoint, index):
         new_layer.conv1.weight.data = old_checkpoint[f"all_modules.{index}.Conv_0.weight"].data
         new_layer.conv1.bias.data = old_checkpoint[f"all_modules.{index}.Conv_0.bias"].data
         new_layer.norm1.weight.data = old_checkpoint[f"all_modules.{index}.GroupNorm_0.weight"].data
@@ -81,35 +79,35 @@ def convert_ncsnpp_checkpoint(checkpoint, config):
     for i, block in enumerate(new_model_architecture.downsample_blocks):
         has_attentions = hasattr(block, "attentions")
         for j in range(len(block.resnets)):
-            set_resnet_weights(block.resnets[j],checkpoint, module_index)
+            set_resnet_weights(block.resnets[j], checkpoint, module_index)
             module_index += 1
             if has_attentions:
-                set_attention_weights(block.attentions[j],checkpoint, module_index)
+                set_attention_weights(block.attentions[j], checkpoint, module_index)
                 module_index += 1
 
         if hasattr(block, "downsamplers") and block.downsamplers is not None:
-            set_resnet_weights(block.resnet_down,checkpoint, module_index)
+            set_resnet_weights(block.resnet_down, checkpoint, module_index)
             module_index += 1
             block.skip_conv.weight.data = checkpoint[f"all_modules.{module_index}.Conv_0.weight"].data
             block.skip_conv.bias.data = checkpoint[f"all_modules.{module_index}.Conv_0.bias"].data
             module_index += 1
 
-
-                    
-    set_resnet_weights(new_model_architecture.mid.resnets[0],checkpoint,module_index)
+    set_resnet_weights(new_model_architecture.mid.resnets[0], checkpoint, module_index)
     module_index += 1
-    set_attention_weights(new_model_architecture.mid.attentions[0],checkpoint, module_index)
+    set_attention_weights(new_model_architecture.mid.attentions[0], checkpoint, module_index)
     module_index += 1
-    set_resnet_weights(new_model_architecture.mid.resnets[1],checkpoint,module_index)
+    set_resnet_weights(new_model_architecture.mid.resnets[1], checkpoint, module_index)
     module_index += 1
 
     for i, block in enumerate(new_model_architecture.upsample_blocks):
         has_attentions = hasattr(block, "attentions")
         for j in range(len(block.resnets)):
-            set_resnet_weights(block.resnets[j],checkpoint, module_index)
+            set_resnet_weights(block.resnets[j], checkpoint, module_index)
             module_index += 1
         if has_attentions:
-            set_attention_weights(block.attentions[0],checkpoint, module_index) # why can there only be a single attention layer for up? 
+            set_attention_weights(
+                block.attentions[0], checkpoint, module_index
+            )  # why can there only be a single attention layer for up?
             module_index += 1
 
         if hasattr(block, "resnet_up") and block.resnet_up is not None:
@@ -119,7 +117,7 @@ def convert_ncsnpp_checkpoint(checkpoint, config):
             block.skip_conv.weight.data = checkpoint[f"all_modules.{module_index}.weight"].data
             block.skip_conv.bias.data = checkpoint[f"all_modules.{module_index}.bias"].data
             module_index += 1
-            set_resnet_weights(block.resnet_up,checkpoint, module_index)
+            set_resnet_weights(block.resnet_up, checkpoint, module_index)
             module_index += 1
 
     new_model_architecture.conv_norm_out.weight.data = checkpoint[f"all_modules.{module_index}.weight"].data
@@ -130,11 +128,16 @@ def convert_ncsnpp_checkpoint(checkpoint, config):
 
     return new_model_architecture.state_dict()
 
+
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
 
     parser.add_argument(
-        "--checkpoint_path", default="/Users/arthurzucker/Work/diffusers/ArthurZ/diffusion_model.pt", type=str, required=False, help="Path to the checkpoint to convert."
+        "--checkpoint_path",
+        default="/Users/arthurzucker/Work/diffusers/ArthurZ/diffusion_model.pt",
+        type=str,
+        required=False,
+        help="Path to the checkpoint to convert.",
     )
 
     parser.add_argument(
@@ -146,19 +149,35 @@ if __name__ == "__main__":
     )
 
     parser.add_argument(
-        "--dump_path", default="/Users/arthurzucker/Work/diffusers/ArthurZ/diffusion_model_new.pt", type=str, required=False, help="Path to the output model."
+        "--dump_path",
+        default="/Users/arthurzucker/Work/diffusers/ArthurZ/diffusion_model_new.pt",
+        type=str,
+        required=False,
+        help="Path to the output model.",
     )
 
     args = parser.parse_args()
 
-
-
-            
     checkpoint = torch.load(args.checkpoint_path, map_location="cpu")
 
     with open(args.config_file) as f:
         config = json.loads(f.read())
 
+    converted_checkpoint = convert_ncsnpp_checkpoint(
+        checkpoint,
+        config,
+    )
+
+    if "sde" in config:
+        del config["sde"]
+
+    model = UNetUnconditionalModel(**config)
+    model.load_state_dict(converted_checkpoint)
+
+    try:
+        scheduler = ScoreSdeVeScheduler.from_config("/".join(args.checkpoint_path.split("/")[:-1]))
 
-    converted_checkpoint = convert_ncsnpp_checkpoint(checkpoint, config,)
-    torch.save(converted_checkpoint, args.dump_path)
+        pipe = ScoreSdeVePipeline(unet=model, scheduler=scheduler)
+        pipe.save_pretrained(args.dump_path)
+    except:
+        model.save_pretrained(args.dump_path)