both worked... but imagedream quality is unmatched

.gitignore

@@ -3,6 +3,6 @@
 **/__pycache__
 *.pyc
 
-weights
+weights*
 models
 sd-v2*

README.md

@@ -3,6 +3,8 @@
 modified from https://github.com/KokeCacao/mvdream-hf.
 
 ### convert weights
+
+MVDream:
 ```bash
 # dependency
 pip install -U omegaconf diffusers safetensors huggingface_hub transformers accelerate
@@ -17,33 +19,20 @@ cd ..
 python convert_mvdream_to_diffusers.py --checkpoint_path models/sd-v2.1-base-4view.pt --dump_path ./weights_mvdream --original_config_file models/sd-v2-base.yaml --half --to_safetensors --test
 ```
 
+ImageDream:
 ```bash
 # download original ckpt
 wget https://huggingface.co/Peng-Wang/ImageDream/resolve/main/sd-v2.1-base-4view-ipmv-local.pt
 wget https://raw.githubusercontent.com/bytedance/ImageDream/main/extern/ImageDream/imagedream/configs/sd_v2_base_ipmv_local.yaml
 
 # convert
-python convert_imagedream_to_diffusers.py --checkpoint_path models/sd-v2.1-base-4view-ipmv-local.pt --dump_path ./weights_imagedream --original_config_file models/sd-v2-base_ipmv_local.yaml --half --to_safetensors --test
+python convert_mvdream_to_diffusers.py --checkpoint_path models/sd-v2.1-base-4view-ipmv-local.pt --dump_path ./weights_imagedream --original_config_file models/sd_v2_base_ipmv_local.yaml --half --to_safetensors --test
 ```
 
 ### usage
 
 example:
 ```bash
-python main.py "a cute owl"
-```
-
-detailed usage:
-```python
-import torch
-import kiui
-from mvdream.pipeline_mvdream import MVDreamPipeline
-
-pipe = MVDreamPipeline.from_pretrained('./weights', torch_dtype=torch.float16)
-pipe = pipe.to("cuda")
-
-prompt = "a photo of an astronaut riding a horse on mars"
-image = pipe(prompt) # np.ndarray [4, 256, 256, 3]
-
-kiui.vis.plot_image(image)
+python run_mvdream.py "a cute owl"
+python run_imagedream.py data/anya_rgba.png
 ```

convert_imagedream_to_diffusers.py

@@ -1,561 +0,0 @@
-# Modified from https://github.com/huggingface/diffusers/blob/bc691231360a4cbc7d19a58742ebb8ed0f05e027/scripts/convert_original_stable_diffusion_to_diffusers.py
-
-import argparse
-import torch
-import sys
-
-sys.path.insert(0, ".")
-
-from diffusers.models import (
-    AutoencoderKL,
-)
-from omegaconf import OmegaConf
-from diffusers.schedulers import DDIMScheduler
-from diffusers.utils import logging
-from typing import Any
-from accelerate import init_empty_weights
-from accelerate.utils import set_module_tensor_to_device
-from imagedream.models import MultiViewUNetModel
-from imagedream.pipeline_imagedream import ImageDreamPipeline
-from transformers import CLIPTextModel, CLIPTokenizer, CLIPVisionModel, CLIPFeatureExtractor
-
-logger = logging.get_logger(__name__)
-
-
-def assign_to_checkpoint(
-    paths,
-    checkpoint,
-    old_checkpoint,
-    attention_paths_to_split=None,
-    additional_replacements=None,
-    config=None,
-):
-    """
-    This does the final conversion step: take locally converted weights and apply a global renaming to them. It splits
-    attention layers, and takes into account additional replacements that may arise.
-    Assigns the weights to the new checkpoint.
-    """
-    assert isinstance(
-        paths, list
-    ), "Paths should be a list of dicts containing 'old' and 'new' keys."
-
-    # Splits the attention layers into three variables.
-    if attention_paths_to_split is not None:
-        for path, path_map in attention_paths_to_split.items():
-            old_tensor = old_checkpoint[path]
-            channels = old_tensor.shape[0] // 3
-
-            target_shape = (-1, channels) if len(old_tensor.shape) == 3 else (-1)
-
-            assert config is not None
-            num_heads = old_tensor.shape[0] // config["num_head_channels"] // 3
-
-            old_tensor = old_tensor.reshape(
-                (num_heads, 3 * channels // num_heads) + old_tensor.shape[1:]
-            )
-            query, key, value = old_tensor.split(channels // num_heads, dim=1)
-
-            checkpoint[path_map["query"]] = query.reshape(target_shape)
-            checkpoint[path_map["key"]] = key.reshape(target_shape)
-            checkpoint[path_map["value"]] = value.reshape(target_shape)
-
-    for path in paths:
-        new_path = path["new"]
-
-        # These have already been assigned
-        if (
-            attention_paths_to_split is not None
-            and new_path in attention_paths_to_split
-        ):
-            continue
-
-        # Global renaming happens here
-        new_path = new_path.replace("middle_block.0", "mid_block.resnets.0")
-        new_path = new_path.replace("middle_block.1", "mid_block.attentions.0")
-        new_path = new_path.replace("middle_block.2", "mid_block.resnets.1")
-
-        if additional_replacements is not None:
-            for replacement in additional_replacements:
-                new_path = new_path.replace(replacement["old"], replacement["new"])
-
-        # proj_attn.weight has to be converted from conv 1D to linear
-        is_attn_weight = "proj_attn.weight" in new_path or (
-            "attentions" in new_path and "to_" in new_path
-        )
-        shape = old_checkpoint[path["old"]].shape
-        if is_attn_weight and len(shape) == 3:
-            checkpoint[new_path] = old_checkpoint[path["old"]][:, :, 0]
-        elif is_attn_weight and len(shape) == 4:
-            checkpoint[new_path] = old_checkpoint[path["old"]][:, :, 0, 0]
-        else:
-            checkpoint[new_path] = old_checkpoint[path["old"]]
-
-
-def shave_segments(path, n_shave_prefix_segments=1):
-    """
-    Removes segments. Positive values shave the first segments, negative shave the last segments.
-    """
-    if n_shave_prefix_segments >= 0:
-        return ".".join(path.split(".")[n_shave_prefix_segments:])
-    else:
-        return ".".join(path.split(".")[:n_shave_prefix_segments])
-
-
-def create_vae_diffusers_config(original_config, image_size: int):
-    """
-    Creates a config for the diffusers based on the config of the LDM model.
-    """
-    vae_params = original_config.model.params.first_stage_config.params.ddconfig
-    _ = original_config.model.params.first_stage_config.params.embed_dim
-
-    block_out_channels = [vae_params.ch * mult for mult in vae_params.ch_mult]
-    down_block_types = ["DownEncoderBlock2D"] * len(block_out_channels)
-    up_block_types = ["UpDecoderBlock2D"] * len(block_out_channels)
-
-    config = {
-        "sample_size": image_size,
-        "in_channels": vae_params.in_channels,
-        "out_channels": vae_params.out_ch,
-        "down_block_types": tuple(down_block_types),
-        "up_block_types": tuple(up_block_types),
-        "block_out_channels": tuple(block_out_channels),
-        "latent_channels": vae_params.z_channels,
-        "layers_per_block": vae_params.num_res_blocks,
-    }
-    return config
-
-
-def convert_ldm_vae_checkpoint(checkpoint, config):
-    # extract state dict for VAE
-    vae_state_dict = {}
-    vae_key = "first_stage_model."
-    keys = list(checkpoint.keys())
-    for key in keys:
-        if key.startswith(vae_key):
-            vae_state_dict[key.replace(vae_key, "")] = checkpoint.get(key)
-
-    new_checkpoint = {}
-
-    new_checkpoint["encoder.conv_in.weight"] = vae_state_dict["encoder.conv_in.weight"]
-    new_checkpoint["encoder.conv_in.bias"] = vae_state_dict["encoder.conv_in.bias"]
-    new_checkpoint["encoder.conv_out.weight"] = vae_state_dict[
-        "encoder.conv_out.weight"
-    ]
-    new_checkpoint["encoder.conv_out.bias"] = vae_state_dict["encoder.conv_out.bias"]
-    new_checkpoint["encoder.conv_norm_out.weight"] = vae_state_dict[
-        "encoder.norm_out.weight"
-    ]
-    new_checkpoint["encoder.conv_norm_out.bias"] = vae_state_dict[
-        "encoder.norm_out.bias"
-    ]
-
-    new_checkpoint["decoder.conv_in.weight"] = vae_state_dict["decoder.conv_in.weight"]
-    new_checkpoint["decoder.conv_in.bias"] = vae_state_dict["decoder.conv_in.bias"]
-    new_checkpoint["decoder.conv_out.weight"] = vae_state_dict[
-        "decoder.conv_out.weight"
-    ]
-    new_checkpoint["decoder.conv_out.bias"] = vae_state_dict["decoder.conv_out.bias"]
-    new_checkpoint["decoder.conv_norm_out.weight"] = vae_state_dict[
-        "decoder.norm_out.weight"
-    ]
-    new_checkpoint["decoder.conv_norm_out.bias"] = vae_state_dict[
-        "decoder.norm_out.bias"
-    ]
-
-    new_checkpoint["quant_conv.weight"] = vae_state_dict["quant_conv.weight"]
-    new_checkpoint["quant_conv.bias"] = vae_state_dict["quant_conv.bias"]
-    new_checkpoint["post_quant_conv.weight"] = vae_state_dict["post_quant_conv.weight"]
-    new_checkpoint["post_quant_conv.bias"] = vae_state_dict["post_quant_conv.bias"]
-
-    # Retrieves the keys for the encoder down blocks only
-    num_down_blocks = len(
-        {
-            ".".join(layer.split(".")[:3])
-            for layer in vae_state_dict
-            if "encoder.down" in layer
-        }
-    )
-    down_blocks = {
-        layer_id: [key for key in vae_state_dict if f"down.{layer_id}" in key]
-        for layer_id in range(num_down_blocks)
-    }
-
-    # Retrieves the keys for the decoder up blocks only
-    num_up_blocks = len(
-        {
-            ".".join(layer.split(".")[:3])
-            for layer in vae_state_dict
-            if "decoder.up" in layer
-        }
-    )
-    up_blocks = {
-        layer_id: [key for key in vae_state_dict if f"up.{layer_id}" in key]
-        for layer_id in range(num_up_blocks)
-    }
-
-    for i in range(num_down_blocks):
-        resnets = [
-            key
-            for key in down_blocks[i]
-            if f"down.{i}" in key and f"down.{i}.downsample" not in key
-        ]
-
-        if f"encoder.down.{i}.downsample.conv.weight" in vae_state_dict:
-            new_checkpoint[
-                f"encoder.down_blocks.{i}.downsamplers.0.conv.weight"
-            ] = vae_state_dict.pop(f"encoder.down.{i}.downsample.conv.weight")
-            new_checkpoint[
-                f"encoder.down_blocks.{i}.downsamplers.0.conv.bias"
-            ] = vae_state_dict.pop(f"encoder.down.{i}.downsample.conv.bias")
-
-        paths = renew_vae_resnet_paths(resnets)
-        meta_path = {"old": f"down.{i}.block", "new": f"down_blocks.{i}.resnets"}
-        assign_to_checkpoint(
-            paths,
-            new_checkpoint,
-            vae_state_dict,
-            additional_replacements=[meta_path],
-            config=config,
-        )
-
-    mid_resnets = [key for key in vae_state_dict if "encoder.mid.block" in key]
-    num_mid_res_blocks = 2
-    for i in range(1, num_mid_res_blocks + 1):
-        resnets = [key for key in mid_resnets if f"encoder.mid.block_{i}" in key]
-
-        paths = renew_vae_resnet_paths(resnets)
-        meta_path = {"old": f"mid.block_{i}", "new": f"mid_block.resnets.{i - 1}"}
-        assign_to_checkpoint(
-            paths,
-            new_checkpoint,
-            vae_state_dict,
-            additional_replacements=[meta_path],
-            config=config,
-        )
-
-    mid_attentions = [key for key in vae_state_dict if "encoder.mid.attn" in key]
-    paths = renew_vae_attention_paths(mid_attentions)
-    meta_path = {"old": "mid.attn_1", "new": "mid_block.attentions.0"}
-    assign_to_checkpoint(
-        paths,
-        new_checkpoint,
-        vae_state_dict,
-        additional_replacements=[meta_path],
-        config=config,
-    )
-    conv_attn_to_linear(new_checkpoint)
-
-    for i in range(num_up_blocks):
-        block_id = num_up_blocks - 1 - i
-        resnets = [
-            key
-            for key in up_blocks[block_id]
-            if f"up.{block_id}" in key and f"up.{block_id}.upsample" not in key
-        ]
-
-        if f"decoder.up.{block_id}.upsample.conv.weight" in vae_state_dict:
-            new_checkpoint[
-                f"decoder.up_blocks.{i}.upsamplers.0.conv.weight"
-            ] = vae_state_dict[f"decoder.up.{block_id}.upsample.conv.weight"]
-            new_checkpoint[
-                f"decoder.up_blocks.{i}.upsamplers.0.conv.bias"
-            ] = vae_state_dict[f"decoder.up.{block_id}.upsample.conv.bias"]
-
-        paths = renew_vae_resnet_paths(resnets)
-        meta_path = {"old": f"up.{block_id}.block", "new": f"up_blocks.{i}.resnets"}
-        assign_to_checkpoint(
-            paths,
-            new_checkpoint,
-            vae_state_dict,
-            additional_replacements=[meta_path],
-            config=config,
-        )
-
-    mid_resnets = [key for key in vae_state_dict if "decoder.mid.block" in key]
-    num_mid_res_blocks = 2
-    for i in range(1, num_mid_res_blocks + 1):
-        resnets = [key for key in mid_resnets if f"decoder.mid.block_{i}" in key]
-
-        paths = renew_vae_resnet_paths(resnets)
-        meta_path = {"old": f"mid.block_{i}", "new": f"mid_block.resnets.{i - 1}"}
-        assign_to_checkpoint(
-            paths,
-            new_checkpoint,
-            vae_state_dict,
-            additional_replacements=[meta_path],
-            config=config,
-        )
-
-    mid_attentions = [key for key in vae_state_dict if "decoder.mid.attn" in key]
-    paths = renew_vae_attention_paths(mid_attentions)
-    meta_path = {"old": "mid.attn_1", "new": "mid_block.attentions.0"}
-    assign_to_checkpoint(
-        paths,
-        new_checkpoint,
-        vae_state_dict,
-        additional_replacements=[meta_path],
-        config=config,
-    )
-    conv_attn_to_linear(new_checkpoint)
-    return new_checkpoint
-
-
-def renew_vae_resnet_paths(old_list, n_shave_prefix_segments=0):
-    """
-    Updates paths inside resnets to the new naming scheme (local renaming)
-    """
-    mapping = []
-    for old_item in old_list:
-        new_item = old_item
-
-        new_item = new_item.replace("nin_shortcut", "conv_shortcut")
-        new_item = shave_segments(
-            new_item, n_shave_prefix_segments=n_shave_prefix_segments
-        )
-
-        mapping.append({"old": old_item, "new": new_item})
-
-    return mapping
-
-
-def renew_vae_attention_paths(old_list, n_shave_prefix_segments=0):
-    """
-    Updates paths inside attentions to the new naming scheme (local renaming)
-    """
-    mapping = []
-    for old_item in old_list:
-        new_item = old_item
-
-        new_item = new_item.replace("norm.weight", "group_norm.weight")
-        new_item = new_item.replace("norm.bias", "group_norm.bias")
-
-        new_item = new_item.replace("q.weight", "to_q.weight")
-        new_item = new_item.replace("q.bias", "to_q.bias")
-
-        new_item = new_item.replace("k.weight", "to_k.weight")
-        new_item = new_item.replace("k.bias", "to_k.bias")
-
-        new_item = new_item.replace("v.weight", "to_v.weight")
-        new_item = new_item.replace("v.bias", "to_v.bias")
-
-        new_item = new_item.replace("proj_out.weight", "to_out.0.weight")
-        new_item = new_item.replace("proj_out.bias", "to_out.0.bias")
-
-        new_item = shave_segments(
-            new_item, n_shave_prefix_segments=n_shave_prefix_segments
-        )
-
-        mapping.append({"old": old_item, "new": new_item})
-
-    return mapping
-
-
-def conv_attn_to_linear(checkpoint):
-    keys = list(checkpoint.keys())
-    attn_keys = ["query.weight", "key.weight", "value.weight"]
-    for key in keys:
-        if ".".join(key.split(".")[-2:]) in attn_keys:
-            if checkpoint[key].ndim > 2:
-                checkpoint[key] = checkpoint[key][:, :, 0, 0]
-        elif "proj_attn.weight" in key:
-            if checkpoint[key].ndim > 2:
-                checkpoint[key] = checkpoint[key][:, :, 0]
-
-
-def create_unet_config(original_config) -> Any:
-    return OmegaConf.to_container(
-        original_config.model.params.unet_config.params, resolve=True
-    )
-
-
-def convert_from_original_imagedream_ckpt(checkpoint_path, original_config_file, device):
-    checkpoint = torch.load(checkpoint_path, map_location=device)
-    # print(f"Checkpoint: {checkpoint.keys()}")
-    torch.cuda.empty_cache()
-
-    original_config = OmegaConf.load(original_config_file)
-    # print(f"Original Config: {original_config}")
-    prediction_type = "epsilon"
-    image_size = 256
-    num_train_timesteps = (
-        getattr(original_config.model.params, "timesteps", None) or 1000
-    )
-    beta_start = getattr(original_config.model.params, "linear_start", None) or 0.02
-    beta_end = getattr(original_config.model.params, "linear_end", None) or 0.085
-    scheduler = DDIMScheduler(
-        beta_end=beta_end,
-        beta_schedule="scaled_linear",
-        beta_start=beta_start,
-        num_train_timesteps=num_train_timesteps,
-        steps_offset=1,
-        clip_sample=False,
-        set_alpha_to_one=False,
-        prediction_type=prediction_type,
-    )
-    scheduler.register_to_config(clip_sample=False)
-
-    # Convert the UNet2DConditionModel model.
-    # upcast_attention = None
-    # unet_config = create_unet_diffusers_config(original_config, image_size=image_size)
-    # unet_config["upcast_attention"] = upcast_attention
-    # with init_empty_weights():
-    #     unet = UNet2DConditionModel(**unet_config)
-    # converted_unet_checkpoint = convert_ldm_unet_checkpoint(
-    #     checkpoint, unet_config, path=None, extract_ema=extract_ema
-    # )
-    # print(f"Unet Config: {original_config.model.params.unet_config.params}")
-    unet_config = create_unet_config(original_config)
-
-    # remove unused configs
-    del unet_config['legacy']
-    del unet_config['use_linear_in_transformer']
-    del unet_config['use_spatial_transformer']
-    del unet_config['ip_mode']
-
-    unet = MultiViewUNetModel(**unet_config)
-    unet.register_to_config(**unet_config)
-    # print(f"Unet State Dict: {unet.state_dict().keys()}")
-    unet.load_state_dict(
-        {
-            key.replace("model.diffusion_model.", ""): value
-            for key, value in checkpoint.items()
-            if key.replace("model.diffusion_model.", "") in unet.state_dict()
-        }
-    )
-    for param_name, param in unet.state_dict().items():
-        set_module_tensor_to_device(unet, param_name, device=device, value=param)
-
-    # Convert the VAE model.
-    vae_config = create_vae_diffusers_config(original_config, image_size=image_size)
-    converted_vae_checkpoint = convert_ldm_vae_checkpoint(checkpoint, vae_config)
-
-    if (
-        "model" in original_config
-        and "params" in original_config.model
-        and "scale_factor" in original_config.model.params
-    ):
-        vae_scaling_factor = original_config.model.params.scale_factor
-    else:
-        vae_scaling_factor = 0.18215  # default SD scaling factor
-
-    vae_config["scaling_factor"] = vae_scaling_factor
-
-    with init_empty_weights():
-        vae = AutoencoderKL(**vae_config)
-
-    for param_name, param in converted_vae_checkpoint.items():
-        set_module_tensor_to_device(vae, param_name, device=device, value=param)
-
-    
-    tokenizer: CLIPTokenizer = CLIPTokenizer.from_pretrained("stabilityai/stable-diffusion-2-1", subfolder="tokenizer")
-    text_encoder: CLIPTextModel = CLIPTextModel.from_pretrained("stabilityai/stable-diffusion-2-1", subfolder="text_encoder").to(device=device)  # type: ignore
-
-    # this is the clip used by sd2.1
-    feature_extractor: CLIPFeatureExtractor = CLIPFeatureExtractor.from_pretrained("laion/CLIP-ViT-H-14-laion2B-s32B-b79K")
-    image_encoder: CLIPVisionModel = CLIPVisionModel.from_pretrained("laion/CLIP-ViT-H-14-laion2B-s32B-b79K")
-    
-    pipe = ImageDreamPipeline(
-        vae=vae,
-        unet=unet,
-        tokenizer=tokenizer,
-        text_encoder=text_encoder,
-        scheduler=scheduler,
-        feature_extractor=feature_extractor,
-        image_encoder=image_encoder,
-    )
-
-    return pipe
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-
-    parser.add_argument(
-        "--checkpoint_path",
-        default=None,
-        type=str,
-        required=True,
-        help="Path to the checkpoint to convert.",
-    )
-    parser.add_argument(
-        "--original_config_file",
-        default=None,
-        type=str,
-        help="The YAML config file corresponding to the original architecture.",
-    )
-    parser.add_argument(
-        "--to_safetensors",
-        action="store_true",
-        help="Whether to store pipeline in safetensors format or not.",
-    )
-    parser.add_argument(
-        "--half", action="store_true", help="Save weights in half precision."
-    )
-    parser.add_argument(
-        "--test",
-        action="store_true",
-        help="Whether to test inference after convertion.",
-    )
-    parser.add_argument(
-        "--dump_path",
-        default=None,
-        type=str,
-        required=True,
-        help="Path to the output model.",
-    )
-    parser.add_argument(
-        "--device", type=str, help="Device to use (e.g. cpu, cuda:0, cuda:1, etc.)"
-    )
-    args = parser.parse_args()
-
-    args.device = torch.device(
-        args.device
-        if args.device is not None
-        else "cuda"
-        if torch.cuda.is_available()
-        else "cpu"
-    )
-
-    pipe = convert_from_original_imagedream_ckpt(
-        checkpoint_path=args.checkpoint_path,
-        original_config_file=args.original_config_file,
-        device=args.device,
-    )
-
-    if args.half:
-        pipe.to(torch_dtype=torch.float16)
-
-    print(f"Saving pipeline to {args.dump_path}...")
-    pipe.save_pretrained(args.dump_path, safe_serialization=args.to_safetensors)
-
-    # TODO: input image...
-    if args.test:
-        try:
-            print(f"Testing each subcomponent of the pipeline...")
-            images = pipe(
-                prompt="Head of Hatsune Miku",
-                negative_prompt="painting, bad quality, flat",
-                output_type="pil",
-                guidance_scale=7.5,
-                num_inference_steps=50,
-                device=args.device,
-            )
-            for i, image in enumerate(images):
-                image.save(f"image_{i}.png")  # type: ignore
-
-            print(f"Testing entire pipeline...")
-            loaded_pipe = ImageDreamPipeline.from_pretrained(args.dump_path, safe_serialization=args.to_safetensors)  # type: ignore
-            images = loaded_pipe(
-                prompt="Head of Hatsune Miku",
-                negative_prompt="painting, bad quality, flat",
-                output_type="pil",
-                guidance_scale=7.5,
-                num_inference_steps=50,
-                device=args.device,
-            )
-            for i, image in enumerate(images):
-                image.save(f"image_{i}.png")  # type: ignore
-        except Exception as e:
-            print(f"Failed to test inference: {e}")
-            raise e from e
-        print("Inference test passed!")

convert_mvdream_to_diffusers.py

@@ -17,7 +17,9 @@ from accelerate import init_empty_weights
 from accelerate.utils import set_module_tensor_to_device
 from mvdream.models import MultiViewUNetModel
 from mvdream.pipeline_mvdream import MVDreamPipeline
-from transformers import CLIPTokenizer, CLIPTextModel
+from transformers import CLIPTextModel, CLIPTokenizer, CLIPVisionModel, CLIPImageProcessor
+
+import kiui
 
 logger = logging.get_logger(__name__)
 
@@ -101,12 +103,20 @@ def shave_segments(path, n_shave_prefix_segments=1):
         return ".".join(path.split(".")[:n_shave_prefix_segments])
 
 
-def create_vae_diffusers_config(original_config, image_size: int):
+def create_vae_diffusers_config(original_config, image_size):
     """
     Creates a config for the diffusers based on the config of the LDM model.
     """
-    vae_params = original_config.model.params.first_stage_config.params.ddconfig
-    _ = original_config.model.params.first_stage_config.params.embed_dim
+
+    
+    if 'imagedream' in original_config.model.target:
+        vae_params = original_config.model.params.vae_config.params.ddconfig
+        _ = original_config.model.params.vae_config.params.embed_dim
+        vae_key = "vae_model."
+    else:
+        vae_params = original_config.model.params.first_stage_config.params.ddconfig
+        _ = original_config.model.params.first_stage_config.params.embed_dim
+        vae_key = "first_stage_model."
 
     block_out_channels = [vae_params.ch * mult for mult in vae_params.ch_mult]
     down_block_types = ["DownEncoderBlock2D"] * len(block_out_channels)
@@ -122,13 +132,12 @@ def create_vae_diffusers_config(original_config, image_size: int):
         "latent_channels": vae_params.z_channels,
         "layers_per_block": vae_params.num_res_blocks,
     }
-    return config
+    return config, vae_key
 
 
-def convert_ldm_vae_checkpoint(checkpoint, config):
+def convert_ldm_vae_checkpoint(checkpoint, config, vae_key):
     # extract state dict for VAE
     vae_state_dict = {}
-    vae_key = "first_stage_model."
     keys = list(checkpoint.keys())
     for key in keys:
         if key.startswith(vae_key):
@@ -394,22 +403,15 @@ def convert_from_original_mvdream_ckpt(checkpoint_path, original_config_file, de
     )
     scheduler.register_to_config(clip_sample=False)
 
-    # Convert the UNet2DConditionModel model.
-    # upcast_attention = None
-    # unet_config = create_unet_diffusers_config(original_config, image_size=image_size)
-    # unet_config["upcast_attention"] = upcast_attention
-    # with init_empty_weights():
-    #     unet = UNet2DConditionModel(**unet_config)
-    # converted_unet_checkpoint = convert_ldm_unet_checkpoint(
-    #     checkpoint, unet_config, path=None, extract_ema=extract_ema
-    # )
-    # print(f"Unet Config: {original_config.model.params.unet_config.params}")
     unet_config = create_unet_config(original_config)
 
     # remove unused configs
-    del unet_config['legacy']
-    del unet_config['use_linear_in_transformer']
-    del unet_config['use_spatial_transformer']
+    unet_config.pop('legacy', None)
+    unet_config.pop('use_linear_in_transformer', None)
+    unet_config.pop('use_spatial_transformer', None)
+    
+    unet_config.pop('ip_mode', None)
+    unet_config.pop('with_ip', None)
 
     unet = MultiViewUNetModel(**unet_config)
     unet.register_to_config(**unet_config)
@@ -425,8 +427,8 @@ def convert_from_original_mvdream_ckpt(checkpoint_path, original_config_file, de
         set_module_tensor_to_device(unet, param_name, device=device, value=param)
 
     # Convert the VAE model.
-    vae_config = create_vae_diffusers_config(original_config, image_size=image_size)
-    converted_vae_checkpoint = convert_ldm_vae_checkpoint(checkpoint, vae_config)
+    vae_config, vae_key = create_vae_diffusers_config(original_config, image_size=image_size)
+    converted_vae_checkpoint = convert_ldm_vae_checkpoint(checkpoint, vae_config, vae_key)
 
     if (
         "model" in original_config
@@ -445,20 +447,17 @@ def convert_from_original_mvdream_ckpt(checkpoint_path, original_config_file, de
     for param_name, param in converted_vae_checkpoint.items():
         set_module_tensor_to_device(vae, param_name, device=device, value=param)
 
-    if original_config.model.params.unet_config.params.context_dim == 768:
-        tokenizer: CLIPTokenizer = CLIPTokenizer.from_pretrained(
-            "openai/clip-vit-large-patch14"
-        )
-        text_encoder: CLIPTextModel = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14").to(device=device)  # type: ignore
-    elif original_config.model.params.unet_config.params.context_dim == 1024:
-        tokenizer: CLIPTokenizer = CLIPTokenizer.from_pretrained(
-            "stabilityai/stable-diffusion-2-1", subfolder="tokenizer"
-        )
-        text_encoder: CLIPTextModel = CLIPTextModel.from_pretrained("stabilityai/stable-diffusion-2-1", subfolder="text_encoder").to(device=device)  # type: ignore
+    # we only supports SD 2.1 based model
+    tokenizer: CLIPTokenizer = CLIPTokenizer.from_pretrained("stabilityai/stable-diffusion-2-1", subfolder="tokenizer")
+    text_encoder: CLIPTextModel = CLIPTextModel.from_pretrained("stabilityai/stable-diffusion-2-1", subfolder="text_encoder").to(device=device)  # type: ignore
+    
+    # imagedream variant
+    if unet.ip_dim > 0:
+        feature_extractor: CLIPImageProcessor = CLIPImageProcessor.from_pretrained("laion/CLIP-ViT-H-14-laion2B-s32B-b79K")
+        image_encoder: CLIPVisionModel = CLIPVisionModel.from_pretrained("laion/CLIP-ViT-H-14-laion2B-s32B-b79K")
     else:
-        raise ValueError(
-            f"Unknown context_dim: {original_config.model.paams.unet_config.params.context_dim}"
-        )
+        feature_extractor = None
+        image_encoder = None
 
     pipe = MVDreamPipeline(
         vae=vae,
@@ -466,6 +465,8 @@ def convert_from_original_mvdream_ckpt(checkpoint_path, original_config_file, de
         tokenizer=tokenizer,
         text_encoder=text_encoder,
         scheduler=scheduler,
+        feature_extractor=feature_extractor,
+        image_encoder=image_encoder,
     )
 
     return pipe
@@ -534,31 +535,63 @@ if __name__ == "__main__":
 
     if args.test:
         try:
-            print(f"Testing each subcomponent of the pipeline...")
-            images = pipe(
-                prompt="Head of Hatsune Miku",
-                negative_prompt="painting, bad quality, flat",
-                output_type="pil",
-                guidance_scale=7.5,
-                num_inference_steps=50,
-                device=args.device,
-            )
-            for i, image in enumerate(images):
-                image.save(f"image_{i}.png")  # type: ignore
-
-            print(f"Testing entire pipeline...")
-            loaded_pipe: MVDreamPipeline = MVDreamPipeline.from_pretrained(args.dump_path, safe_serialization=args.to_safetensors)  # type: ignore
-            images = loaded_pipe(
-                prompt="Head of Hatsune Miku",
-                negative_prompt="painting, bad quality, flat",
-                output_type="pil",
-                guidance_scale=7.5,
-                num_inference_steps=50,
-                device=args.device,
-            )
-            for i, image in enumerate(images):
-                image.save(f"image_{i}.png")  # type: ignore
+            # mvdream
+            if pipe.unet.ip_dim == 0:
+                print(f"Testing each subcomponent of the pipeline...")
+                images = pipe(
+                    prompt="Head of Hatsune Miku",
+                    negative_prompt="painting, bad quality, flat",
+                    output_type="pil",
+                    guidance_scale=7.5,
+                    num_inference_steps=50,
+                    device=args.device,
+                )
+                for i, image in enumerate(images):
+                    image.save(f"test_image_{i}.png")  # type: ignore
+
+                print(f"Testing entire pipeline...")
+                loaded_pipe = MVDreamPipeline.from_pretrained(args.dump_path)  # type: ignore
+                images = loaded_pipe(
+                    prompt="Head of Hatsune Miku",
+                    negative_prompt="painting, bad quality, flat",
+                    output_type="pil",
+                    guidance_scale=7.5,
+                    num_inference_steps=50,
+                    device=args.device,
+                )
+                for i, image in enumerate(images):
+                    image.save(f"test_image_{i}.png")  # type: ignore
+            # imagedream
+            else:
+                input_image = kiui.read_image('data/anya_rgba.png', mode='float')
+                print(f"Testing each subcomponent of the pipeline...")
+                images = pipe(
+                    image=input_image,
+                    prompt="",
+                    negative_prompt="painting, bad quality, flat",
+                    output_type="pil",
+                    guidance_scale=5.0,
+                    num_inference_steps=50,
+                    device=args.device,
+                )
+                for i, image in enumerate(images):
+                    image.save(f"test_image_{i}.png")  # type: ignore
+
+                print(f"Testing entire pipeline...")
+                loaded_pipe = MVDreamPipeline.from_pretrained(args.dump_path)  # type: ignore
+                images = loaded_pipe(
+                    image=input_image,
+                    prompt="",
+                    negative_prompt="painting, bad quality, flat",
+                    output_type="pil",
+                    guidance_scale=5.0,
+                    num_inference_steps=50,
+                    device=args.device,
+                )
+                for i, image in enumerate(images):
+                    image.save(f"test_image_{i}.png")  # type: ignore
+                
+
+            print("Inference test passed!")
         except Exception as e:
             print(f"Failed to test inference: {e}")
-            raise e from e
-        print("Inference test passed!")

imagedream/attention.py

@@ -1,259 +0,0 @@
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-from inspect import isfunction
-from einops import rearrange, repeat
-from typing import Optional, Any
-
-# require xformers
-import xformers  # type: ignore
-import xformers.ops  # type: ignore
-
-from .util import checkpoint, zero_module
-
-def default(val, d):
-    if val is not None:
-        return val
-    return d() if isfunction(d) else d
-
-
-class GEGLU(nn.Module):
-    def __init__(self, dim_in, dim_out):
-        super().__init__()
-        self.proj = nn.Linear(dim_in, dim_out * 2)
-
-    def forward(self, x):
-        x, gate = self.proj(x).chunk(2, dim=-1)
-        return x * F.gelu(gate)
-
-
-class FeedForward(nn.Module):
-    def __init__(self, dim, dim_out=None, mult=4, glu=False, dropout=0.0):
-        super().__init__()
-        inner_dim = int(dim * mult)
-        dim_out = default(dim_out, dim)
-        project_in = (
-            nn.Sequential(nn.Linear(dim, inner_dim), nn.GELU())
-            if not glu
-            else GEGLU(dim, inner_dim)
-        )
-
-        self.net = nn.Sequential(
-            project_in, nn.Dropout(dropout), nn.Linear(inner_dim, dim_out)
-        )
-
-    def forward(self, x):
-        return self.net(x)
-
-
-class MemoryEfficientCrossAttention(nn.Module):
-    # https://github.com/MatthieuTPHR/diffusers/blob/d80b531ff8060ec1ea982b65a1b8df70f73aa67c/src/diffusers/models/attention.py#L223
-    def __init__(
-            self, 
-            query_dim, 
-            context_dim=None, 
-            heads=8, 
-            dim_head=64, 
-            dropout=0.0,
-            with_ip=False,
-            ip_dim=16,
-            ip_weight=1,
-        ):
-        super().__init__()
-        
-        inner_dim = dim_head * heads
-        context_dim = default(context_dim, query_dim)
-
-        self.heads = heads
-        self.dim_head = dim_head
-
-        self.with_ip = with_ip and (context_dim is not None)
-        self.ip_dim = ip_dim
-        self.ip_weight = ip_weight
-
-        if self.with_ip:
-            self.to_k_ip = nn.Linear(context_dim, inner_dim, bias=False)
-            self.to_v_ip = nn.Linear(context_dim, inner_dim, bias=False)
-
-        self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
-        self.to_k = nn.Linear(context_dim, inner_dim, bias=False)
-        self.to_v = nn.Linear(context_dim, inner_dim, bias=False)
-
-        self.to_out = nn.Sequential(
-            nn.Linear(inner_dim, query_dim), nn.Dropout(dropout)
-        )
-        self.attention_op: Optional[Any] = None
-
-    def forward(self, x, context=None):
-        q = self.to_q(x)
-        context = default(context, x)
-
-        if self.with_ip:
-            # context dim [(b frame_num), (77 + img_token), 1024]
-            token_len = context.shape[1]
-            context_ip = context[:, -self.ip_dim :, :]
-            k_ip = self.to_k_ip(context_ip)
-            v_ip = self.to_v_ip(context_ip)
-            context = context[:, : (token_len - self.ip_dim), :]
-
-        k = self.to_k(context)
-        v = self.to_v(context)
-
-        b, _, _ = q.shape
-        q, k, v = map(
-            lambda t: t.unsqueeze(3)
-            .reshape(b, t.shape[1], self.heads, self.dim_head)
-            .permute(0, 2, 1, 3)
-            .reshape(b * self.heads, t.shape[1], self.dim_head)
-            .contiguous(),
-            (q, k, v),
-        )
-
-        # actually compute the attention, what we cannot get enough of
-        out = xformers.ops.memory_efficient_attention(
-            q, k, v, attn_bias=None, op=self.attention_op
-        )
-
-        if self.with_ip:
-            k_ip, v_ip = map(
-                lambda t: t.unsqueeze(3)
-                .reshape(b, t.shape[1], self.heads, self.dim_head)
-                .permute(0, 2, 1, 3)
-                .reshape(b * self.heads, t.shape[1], self.dim_head)
-                .contiguous(),
-                (k_ip, v_ip),
-            )
-            # actually compute the attention, what we cannot get enough of
-            out_ip = xformers.ops.memory_efficient_attention(
-                q, k_ip, v_ip, attn_bias=None, op=self.attention_op
-            )
-            out = out + self.ip_weight * out_ip
-
-        out = (
-            out.unsqueeze(0)
-            .reshape(b, self.heads, out.shape[1], self.dim_head)
-            .permute(0, 2, 1, 3)
-            .reshape(b, out.shape[1], self.heads * self.dim_head)
-        )
-        return self.to_out(out)
-
-
-class BasicTransformerBlock3D(nn.Module):
-    
-    def __init__(
-        self,
-        dim,
-        context_dim,
-        n_heads,
-        d_head,
-        dropout=0.0,
-        gated_ff=True,
-        checkpoint=True,
-        with_ip=False,
-        ip_dim=16,
-        ip_weight=1,
-    ):
-        super().__init__()
-
-        self.attn1 = MemoryEfficientCrossAttention(
-            query_dim=dim,
-            context_dim=None, # self-attention
-            heads=n_heads,
-            dim_head=d_head,
-            dropout=dropout,
-        )
-        self.ff = FeedForward(dim, dropout=dropout, glu=gated_ff)
-        self.attn2 = MemoryEfficientCrossAttention(
-            query_dim=dim,
-            context_dim=context_dim,
-            heads=n_heads,
-            dim_head=d_head,
-            dropout=dropout,
-            # ip only applies to cross-attention
-            with_ip=with_ip,
-            ip_dim=ip_dim,
-            ip_weight=ip_weight,
-        ) 
-        self.norm1 = nn.LayerNorm(dim)
-        self.norm2 = nn.LayerNorm(dim)
-        self.norm3 = nn.LayerNorm(dim)
-        self.checkpoint = checkpoint
-
-    def forward(self, x, context=None, num_frames=1):
-        return checkpoint(
-            self._forward, (x, context, num_frames), self.parameters(), self.checkpoint
-        )
-
-    def _forward(self, x, context=None, num_frames=1):
-        x = rearrange(x, "(b f) l c -> b (f l) c", f=num_frames).contiguous()
-        x = self.attn1(self.norm1(x), context=None) + x
-        x = rearrange(x, "b (f l) c -> (b f) l c", f=num_frames).contiguous()
-        x = self.attn2(self.norm2(x), context=context) + x
-        x = self.ff(self.norm3(x)) + x
-        return x
-
-
-class SpatialTransformer3D(nn.Module):
-
-    def __init__(
-        self,
-        in_channels,
-        n_heads,
-        d_head,
-        context_dim, # cross attention input dim
-        depth=1,
-        dropout=0.0,
-        with_ip=False,
-        ip_dim=16,
-        ip_weight=1,
-        use_checkpoint=True,
-    ):
-        super().__init__()
-
-        if not isinstance(context_dim, list):
-            context_dim = [context_dim]
-
-        self.in_channels = in_channels
-
-        inner_dim = n_heads * d_head
-        self.norm = nn.GroupNorm(
-            num_groups=32, num_channels=in_channels, eps=1e-6, affine=True
-        )
-        self.proj_in = nn.Linear(in_channels, inner_dim)
-
-        self.transformer_blocks = nn.ModuleList(
-            [
-                BasicTransformerBlock3D(
-                    inner_dim,
-                    n_heads,
-                    d_head,
-                    context_dim=context_dim[d],
-                    dropout=dropout,
-                    checkpoint=use_checkpoint,
-                    with_ip=with_ip,
-                    ip_dim=ip_dim,
-                    ip_weight=ip_weight,
-                )
-                for d in range(depth)
-            ]
-        )
-        
-        self.proj_out = zero_module(nn.Linear(in_channels, inner_dim))
-        
-
-    def forward(self, x, context=None, num_frames=1):
-        # note: if no context is given, cross-attention defaults to self-attention
-        if not isinstance(context, list):
-            context = [context]
-        b, c, h, w = x.shape
-        x_in = x
-        x = self.norm(x)
-        x = rearrange(x, "b c h w -> b (h w) c").contiguous()
-        x = self.proj_in(x)
-        for i, block in enumerate(self.transformer_blocks):
-            x = block(x, context=context[i], num_frames=num_frames)
-        x = self.proj_out(x)
-        x = rearrange(x, "b (h w) c -> b c h w", h=h, w=w).contiguous()
-        
-        return x + x_in

imagedream/models.py

@@ -1,627 +0,0 @@
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from diffusers.configuration_utils import ConfigMixin
-from diffusers.models.modeling_utils import ModelMixin
-from typing import Any, List, Optional
-from torch import Tensor
-
-from .util import (
-    checkpoint,
-    conv_nd,
-    avg_pool_nd,
-    zero_module,
-    timestep_embedding,
-)
-from .attention import SpatialTransformer3D
-from .adaptor import Resampler, ImageProjModel
-
-class CondSequential(nn.Sequential):
-    """
-    A sequential module that passes timestep embeddings to the children that
-    support it as an extra input.
-    """
-
-    def forward(self, x, emb, context=None, num_frames=1):
-        for layer in self:
-            if isinstance(layer, ResBlock):
-                x = layer(x, emb)
-            elif isinstance(layer, SpatialTransformer3D):
-                x = layer(x, context, num_frames=num_frames)
-            else:
-                x = layer(x)
-        return x
-
-
-class Upsample(nn.Module):
-    """
-    An upsampling layer with an optional convolution.
-    :param channels: channels in the inputs and outputs.
-    :param use_conv: a bool determining if a convolution is applied.
-    :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then
-                 upsampling occurs in the inner-two dimensions.
-    """
-
-    def __init__(self, channels, use_conv, dims=2, out_channels=None, padding=1):
-        super().__init__()
-        self.channels = channels
-        self.out_channels = out_channels or channels
-        self.use_conv = use_conv
-        self.dims = dims
-        if use_conv:
-            self.conv = conv_nd(
-                dims, self.channels, self.out_channels, 3, padding=padding
-            )
-
-    def forward(self, x):
-        assert x.shape[1] == self.channels
-        if self.dims == 3:
-            x = F.interpolate(
-                x, (x.shape[2], x.shape[3] * 2, x.shape[4] * 2), mode="nearest"
-            )
-        else:
-            x = F.interpolate(x, scale_factor=2, mode="nearest")
-        if self.use_conv:
-            x = self.conv(x)
-        return x
-
-
-class Downsample(nn.Module):
-    """
-    A downsampling layer with an optional convolution.
-    :param channels: channels in the inputs and outputs.
-    :param use_conv: a bool determining if a convolution is applied.
-    :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then
-                 downsampling occurs in the inner-two dimensions.
-    """
-
-    def __init__(self, channels, use_conv, dims=2, out_channels=None, padding=1):
-        super().__init__()
-        self.channels = channels
-        self.out_channels = out_channels or channels
-        self.use_conv = use_conv
-        self.dims = dims
-        stride = 2 if dims != 3 else (1, 2, 2)
-        if use_conv:
-            self.op = conv_nd(
-                dims,
-                self.channels,
-                self.out_channels,
-                3,
-                stride=stride,
-                padding=padding,
-            )
-        else:
-            assert self.channels == self.out_channels
-            self.op = avg_pool_nd(dims, kernel_size=stride, stride=stride)
-
-    def forward(self, x):
-        assert x.shape[1] == self.channels
-        return self.op(x)
-
-
-class ResBlock(nn.Module):
-    """
-    A residual block that can optionally change the number of channels.
-    :param channels: the number of input channels.
-    :param emb_channels: the number of timestep embedding channels.
-    :param dropout: the rate of dropout.
-    :param out_channels: if specified, the number of out channels.
-    :param use_conv: if True and out_channels is specified, use a spatial
-        convolution instead of a smaller 1x1 convolution to change the
-        channels in the skip connection.
-    :param dims: determines if the signal is 1D, 2D, or 3D.
-    :param use_checkpoint: if True, use gradient checkpointing on this module.
-    :param up: if True, use this block for upsampling.
-    :param down: if True, use this block for downsampling.
-    """
-
-    def __init__(
-        self,
-        channels,
-        emb_channels,
-        dropout,
-        out_channels=None,
-        use_conv=False,
-        use_scale_shift_norm=False,
-        dims=2,
-        use_checkpoint=False,
-        up=False,
-        down=False,
-    ):
-        super().__init__()
-        self.channels = channels
-        self.emb_channels = emb_channels
-        self.dropout = dropout
-        self.out_channels = out_channels or channels
-        self.use_conv = use_conv
-        self.use_checkpoint = use_checkpoint
-        self.use_scale_shift_norm = use_scale_shift_norm
-
-        self.in_layers = nn.Sequential(
-            nn.GroupNorm(32, channels),
-            nn.SiLU(),
-            conv_nd(dims, channels, self.out_channels, 3, padding=1),
-        )
-
-        self.updown = up or down
-
-        if up:
-            self.h_upd = Upsample(channels, False, dims)
-            self.x_upd = Upsample(channels, False, dims)
-        elif down:
-            self.h_upd = Downsample(channels, False, dims)
-            self.x_upd = Downsample(channels, False, dims)
-        else:
-            self.h_upd = self.x_upd = nn.Identity()
-
-        self.emb_layers = nn.Sequential(
-            nn.SiLU(),
-            nn.Linear(
-                emb_channels,
-                2 * self.out_channels if use_scale_shift_norm else self.out_channels,
-            ),
-        )
-        self.out_layers = nn.Sequential(
-            nn.GroupNorm(32, self.out_channels),
-            nn.SiLU(),
-            nn.Dropout(p=dropout),
-            zero_module(
-                conv_nd(dims, self.out_channels, self.out_channels, 3, padding=1)
-            ),
-        )
-
-        if self.out_channels == channels:
-            self.skip_connection = nn.Identity()
-        elif use_conv:
-            self.skip_connection = conv_nd(
-                dims, channels, self.out_channels, 3, padding=1
-            )
-        else:
-            self.skip_connection = conv_nd(dims, channels, self.out_channels, 1)
-
-    def forward(self, x, emb):
-        """
-        Apply the block to a Tensor, conditioned on a timestep embedding.
-        :param x: an [N x C x ...] Tensor of features.
-        :param emb: an [N x emb_channels] Tensor of timestep embeddings.
-        :return: an [N x C x ...] Tensor of outputs.
-        """
-        return checkpoint(
-            self._forward, (x, emb), self.parameters(), self.use_checkpoint
-        )
-
-    def _forward(self, x, emb):
-        if self.updown:
-            in_rest, in_conv = self.in_layers[:-1], self.in_layers[-1]
-            h = in_rest(x)
-            h = self.h_upd(h)
-            x = self.x_upd(x)
-            h = in_conv(h)
-        else:
-            h = self.in_layers(x)
-        emb_out = self.emb_layers(emb).type(h.dtype)
-        while len(emb_out.shape) < len(h.shape):
-            emb_out = emb_out[..., None]
-        if self.use_scale_shift_norm:
-            out_norm, out_rest = self.out_layers[0], self.out_layers[1:]
-            scale, shift = torch.chunk(emb_out, 2, dim=1)
-            h = out_norm(h) * (1 + scale) + shift
-            h = out_rest(h)
-        else:
-            h = h + emb_out
-            h = self.out_layers(h)
-        return self.skip_connection(x) + h
-
-
-class MultiViewUNetModel(ModelMixin, ConfigMixin):
-    """
-    The full multi-view UNet model with attention, timestep embedding and camera embedding.
-    :param in_channels: channels in the input Tensor.
-    :param model_channels: base channel count for the model.
-    :param out_channels: channels in the output Tensor.
-    :param num_res_blocks: number of residual blocks per downsample.
-    :param attention_resolutions: a collection of downsample rates at which
-        attention will take place. May be a set, list, or tuple.
-        For example, if this contains 4, then at 4x downsampling, attention
-        will be used.
-    :param dropout: the dropout probability.
-    :param channel_mult: channel multiplier for each level of the UNet.
-    :param conv_resample: if True, use learned convolutions for upsampling and
-        downsampling.
-    :param dims: determines if the signal is 1D, 2D, or 3D.
-    :param num_classes: if specified (as an int), then this model will be
-        class-conditional with `num_classes` classes.
-    :param use_checkpoint: use gradient checkpointing to reduce memory usage.
-    :param num_heads: the number of attention heads in each attention layer.
-    :param num_heads_channels: if specified, ignore num_heads and instead use
-                               a fixed channel width per attention head.
-    :param num_heads_upsample: works with num_heads to set a different number
-                               of heads for upsampling. Deprecated.
-    :param use_scale_shift_norm: use a FiLM-like conditioning mechanism.
-    :param resblock_updown: use residual blocks for up/downsampling.
-    :param use_new_attention_order: use a different attention pattern for potentially
-                                    increased efficiency.
-    :param camera_dim: dimensionality of camera input.
-    """
-
-    def __init__(
-        self,
-        image_size,
-        in_channels,
-        model_channels,
-        out_channels,
-        num_res_blocks,
-        attention_resolutions,
-        dropout=0,
-        channel_mult=(1, 2, 4, 8),
-        conv_resample=True,
-        dims=2,
-        num_classes=None,
-        use_checkpoint=False,
-        num_heads=-1,
-        num_head_channels=-1,
-        num_heads_upsample=-1,
-        use_scale_shift_norm=False,
-        resblock_updown=False,
-        transformer_depth=1,  # custom transformer support
-        context_dim=None,  # custom transformer support
-        n_embed=None,  # custom support for prediction of discrete ids into codebook of first stage vq model
-        disable_self_attentions=None,
-        num_attention_blocks=None,
-        disable_middle_self_attn=False,
-        adm_in_channels=None,
-        camera_dim=None,
-        with_ip=True,
-        ip_dim=16,
-        ip_weight=1.0,
-        **kwargs,
-    ):
-        super().__init__()
-        assert context_dim is not None
-        
-        if num_heads_upsample == -1:
-            num_heads_upsample = num_heads
-
-        if num_heads == -1:
-            assert (
-                num_head_channels != -1
-            ), "Either num_heads or num_head_channels has to be set"
-
-        if num_head_channels == -1:
-            assert (
-                num_heads != -1
-            ), "Either num_heads or num_head_channels has to be set"
-
-        self.image_size = image_size
-        self.in_channels = in_channels
-        self.model_channels = model_channels
-        self.out_channels = out_channels
-        if isinstance(num_res_blocks, int):
-            self.num_res_blocks = len(channel_mult) * [num_res_blocks]
-        else:
-            if len(num_res_blocks) != len(channel_mult):
-                raise ValueError(
-                    "provide num_res_blocks either as an int (globally constant) or "
-                    "as a list/tuple (per-level) with the same length as channel_mult"
-                )
-            self.num_res_blocks = num_res_blocks
-        
-        if num_attention_blocks is not None:
-            assert len(num_attention_blocks) == len(self.num_res_blocks)
-            assert all(
-                map(
-                    lambda i: self.num_res_blocks[i] >= num_attention_blocks[i],
-                    range(len(num_attention_blocks)),
-                )
-            )
-            print(
-                f"Constructor of UNetModel received num_attention_blocks={num_attention_blocks}. "
-                f"This option has LESS priority than attention_resolutions {attention_resolutions}, "
-                f"i.e., in cases where num_attention_blocks[i] > 0 but 2**i not in attention_resolutions, "
-                f"attention will still not be set."
-            )
-
-        self.attention_resolutions = attention_resolutions
-        self.dropout = dropout
-        self.channel_mult = channel_mult
-        self.conv_resample = conv_resample
-        self.num_classes = num_classes
-        self.use_checkpoint = use_checkpoint
-        self.num_heads = num_heads
-        self.num_head_channels = num_head_channels
-        self.num_heads_upsample = num_heads_upsample
-        self.predict_codebook_ids = n_embed is not None
-
-        self.with_ip = with_ip
-        self.ip_dim = ip_dim
-        self.ip_weight = ip_weight
-
-        if self.with_ip and self.ip_dim > 0:
-            self.image_embed = Resampler(
-                dim=context_dim,
-                depth=4,
-                dim_head=64,
-                heads=12,
-                num_queries=ip_dim,  # num token
-                embedding_dim=1280,
-                output_dim=context_dim,
-                ff_mult=4,
-            )
-
-        time_embed_dim = model_channels * 4
-        self.time_embed = nn.Sequential(
-            nn.Linear(model_channels, time_embed_dim),
-            nn.SiLU(),
-            nn.Linear(time_embed_dim, time_embed_dim),
-        )
-
-        if camera_dim is not None:
-            time_embed_dim = model_channels * 4
-            self.camera_embed = nn.Sequential(
-                nn.Linear(camera_dim, time_embed_dim),
-                nn.SiLU(),
-                nn.Linear(time_embed_dim, time_embed_dim),
-            )
-
-        if self.num_classes is not None:
-            if isinstance(self.num_classes, int):
-                self.label_emb = nn.Embedding(self.num_classes, time_embed_dim)
-            elif self.num_classes == "continuous":
-                # print("setting up linear c_adm embedding layer")
-                self.label_emb = nn.Linear(1, time_embed_dim)
-            elif self.num_classes == "sequential":
-                assert adm_in_channels is not None
-                self.label_emb = nn.Sequential(
-                    nn.Sequential(
-                        nn.Linear(adm_in_channels, time_embed_dim),
-                        nn.SiLU(),
-                        nn.Linear(time_embed_dim, time_embed_dim),
-                    )
-                )
-            else:
-                raise ValueError()
-
-        self.input_blocks = nn.ModuleList(
-            [
-                CondSequential(
-                    conv_nd(dims, in_channels, model_channels, 3, padding=1)
-                )
-            ]
-        )
-        self._feature_size = model_channels
-        input_block_chans = [model_channels]
-        ch = model_channels
-        ds = 1
-        for level, mult in enumerate(channel_mult):
-            for nr in range(self.num_res_blocks[level]):
-                layers: List[Any] = [
-                    ResBlock(
-                        ch,
-                        time_embed_dim,
-                        dropout,
-                        out_channels=mult * model_channels,
-                        dims=dims,
-                        use_checkpoint=use_checkpoint,
-                        use_scale_shift_norm=use_scale_shift_norm,
-                    )
-                ]
-                ch = mult * model_channels
-                if ds in attention_resolutions:
-                    if num_head_channels == -1:
-                        dim_head = ch // num_heads
-                    else:
-                        num_heads = ch // num_head_channels
-                        dim_head = num_head_channels
-
-                    if num_attention_blocks is None or nr < num_attention_blocks[level]:
-                        layers.append(
-                            SpatialTransformer3D(
-                                ch,
-                                num_heads,
-                                dim_head,
-                                context_dim=context_dim,
-                                depth=transformer_depth,
-                                use_checkpoint=use_checkpoint,
-                                with_ip=self.with_ip,
-                                ip_dim=self.ip_dim,
-                                ip_weight=self.ip_weight,
-                            )
-                        )
-                self.input_blocks.append(CondSequential(*layers))
-                self._feature_size += ch
-                input_block_chans.append(ch)
-            if level != len(channel_mult) - 1:
-                out_ch = ch
-                self.input_blocks.append(
-                    CondSequential(
-                        ResBlock(
-                            ch,
-                            time_embed_dim,
-                            dropout,
-                            out_channels=out_ch,
-                            dims=dims,
-                            use_checkpoint=use_checkpoint,
-                            use_scale_shift_norm=use_scale_shift_norm,
-                            down=True,
-                        )
-                        if resblock_updown
-                        else Downsample(
-                            ch, conv_resample, dims=dims, out_channels=out_ch
-                        )
-                    )
-                )
-                ch = out_ch
-                input_block_chans.append(ch)
-                ds *= 2
-                self._feature_size += ch
-
-        if num_head_channels == -1:
-            dim_head = ch // num_heads
-        else:
-            num_heads = ch // num_head_channels
-            dim_head = num_head_channels
-        
-        self.middle_block = CondSequential(
-            ResBlock(
-                ch,
-                time_embed_dim,
-                dropout,
-                dims=dims,
-                use_checkpoint=use_checkpoint,
-                use_scale_shift_norm=use_scale_shift_norm,
-            ),
-            SpatialTransformer3D(
-                ch,
-                num_heads,
-                dim_head,
-                context_dim=context_dim,
-                depth=transformer_depth,
-                use_checkpoint=use_checkpoint,
-                with_ip=self.with_ip,
-                ip_dim=self.ip_dim,
-                ip_weight=self.ip_weight,
-            ), 
-            ResBlock(
-                ch,
-                time_embed_dim,
-                dropout,
-                dims=dims,
-                use_checkpoint=use_checkpoint,
-                use_scale_shift_norm=use_scale_shift_norm,
-            ),
-        )
-        self._feature_size += ch
-
-        self.output_blocks = nn.ModuleList([])
-        for level, mult in list(enumerate(channel_mult))[::-1]:
-            for i in range(self.num_res_blocks[level] + 1):
-                ich = input_block_chans.pop()
-                layers = [
-                    ResBlock(
-                        ch + ich,
-                        time_embed_dim,
-                        dropout,
-                        out_channels=model_channels * mult,
-                        dims=dims,
-                        use_checkpoint=use_checkpoint,
-                        use_scale_shift_norm=use_scale_shift_norm,
-                    )
-                ]
-                ch = model_channels * mult
-                if ds in attention_resolutions:
-                    if num_head_channels == -1:
-                        dim_head = ch // num_heads
-                    else:
-                        num_heads = ch // num_head_channels
-                        dim_head = num_head_channels
-
-                    if num_attention_blocks is None or i < num_attention_blocks[level]:
-                        layers.append(
-                            SpatialTransformer3D(
-                                ch,
-                                num_heads,
-                                dim_head,
-                                context_dim=context_dim,
-                                depth=transformer_depth,
-                                use_checkpoint=use_checkpoint,
-                                with_ip=self.with_ip,
-                                ip_dim=self.ip_dim,
-                                ip_weight=self.ip_weight,
-                            )
-                        )
-                if level and i == self.num_res_blocks[level]:
-                    out_ch = ch
-                    layers.append(
-                        ResBlock(
-                            ch,
-                            time_embed_dim,
-                            dropout,
-                            out_channels=out_ch,
-                            dims=dims,
-                            use_checkpoint=use_checkpoint,
-                            use_scale_shift_norm=use_scale_shift_norm,
-                            up=True,
-                        )
-                        if resblock_updown
-                        else Upsample(ch, conv_resample, dims=dims, out_channels=out_ch)
-                    )
-                    ds //= 2
-                self.output_blocks.append(CondSequential(*layers))
-                self._feature_size += ch
-
-        self.out = nn.Sequential(
-            nn.GroupNorm(32, ch),
-            nn.SiLU(),
-            zero_module(conv_nd(dims, model_channels, out_channels, 3, padding=1)),
-        )
-        if self.predict_codebook_ids:
-            self.id_predictor = nn.Sequential(
-                nn.GroupNorm(32, ch),
-                conv_nd(dims, model_channels, n_embed, 1),
-                # nn.LogSoftmax(dim=1)  # change to cross_entropy and produce non-normalized logits
-            )
-
-    def forward(
-        self,
-        x,
-        timesteps=None,
-        context=None,
-        y: Optional[Tensor] = None,
-        camera=None,
-        num_frames=1,
-        # should be provided if with_ip
-        ip = None,
-        ip_img = None,
-        **kwargs,
-    ):
-        """
-        Apply the model to an input batch.
-        :param x: an [(N x F) x C x ...] Tensor of inputs. F is the number of frames (views).
-        :param timesteps: a 1-D batch of timesteps.
-        :param context: conditioning plugged in via crossattn
-        :param y: an [N] Tensor of labels, if class-conditional.
-        :param num_frames: a integer indicating number of frames for tensor reshaping.
-        :return: an [(N x F) x C x ...] Tensor of outputs. F is the number of frames (views).
-        """
-        assert (
-            x.shape[0] % num_frames == 0
-        ), "[UNet] input batch size must be dividable by num_frames!"
-        assert (y is not None) == (
-            self.num_classes is not None
-        ), "must specify y if and only if the model is class-conditional"
-        hs = []
-        t_emb = timestep_embedding(
-            timesteps, self.model_channels, repeat_only=False
-        ).to(x.dtype)
-
-        emb = self.time_embed(t_emb)
-
-        if self.num_classes is not None:
-            assert y is not None
-            assert y.shape[0] == x.shape[0]
-            emb = emb + self.label_emb(y)
-
-        # Add camera embeddings
-        if camera is not None:
-            assert camera.shape[0] == emb.shape[0]
-            emb = emb + self.camera_embed(camera)
-        
-        if self.with_ip:
-            x[(num_frames - 1) :: num_frames, :, :, :] = ip_img
-            ip_emb = self.image_embed(ip)
-            context = torch.cat((context, ip_emb), 1)
-
-        h = x
-        for module in self.input_blocks:
-            h = module(h, emb, context, num_frames=num_frames)
-            hs.append(h)
-        h = self.middle_block(h, emb, context, num_frames=num_frames)
-        for module in self.output_blocks:
-            h = torch.cat([h, hs.pop()], dim=1)
-            h = module(h, emb, context, num_frames=num_frames)
-        h = h.type(x.dtype)
-        if self.predict_codebook_ids:
-            return self.id_predictor(h)
-        else:
-            return self.out(h)

imagedream/pipeline_imagedream.py

@@ -1,620 +0,0 @@
-import torch
-import inspect
-import numpy as np
-from typing import Callable, List, Optional, Union
-from transformers import CLIPTextModel, CLIPTokenizer, CLIPVisionModel, CLIPFeatureExtractor
-from diffusers import AutoencoderKL, DiffusionPipeline
-from diffusers.utils import (
-    deprecate,
-    is_accelerate_available,
-    is_accelerate_version,
-    logging,
-)
-from diffusers.configuration_utils import FrozenDict
-from diffusers.schedulers import DDIMScheduler
-from diffusers.utils.torch_utils import randn_tensor
-
-from .models import MultiViewUNetModel
-
-import kiui
-
-logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
-
-
-def create_camera_to_world_matrix(elevation, azimuth):
-    elevation = np.radians(elevation)
-    azimuth = np.radians(azimuth)
-    # Convert elevation and azimuth angles to Cartesian coordinates on a unit sphere
-    x = np.cos(elevation) * np.sin(azimuth)
-    y = np.sin(elevation)
-    z = np.cos(elevation) * np.cos(azimuth)
-
-    # Calculate camera position, target, and up vectors
-    camera_pos = np.array([x, y, z])
-    target = np.array([0, 0, 0])
-    up = np.array([0, 1, 0])
-
-    # Construct view matrix
-    forward = target - camera_pos
-    forward /= np.linalg.norm(forward)
-    right = np.cross(forward, up)
-    right /= np.linalg.norm(right)
-    new_up = np.cross(right, forward)
-    new_up /= np.linalg.norm(new_up)
-    cam2world = np.eye(4)
-    cam2world[:3, :3] = np.array([right, new_up, -forward]).T
-    cam2world[:3, 3] = camera_pos
-    return cam2world
-
-
-def convert_opengl_to_blender(camera_matrix):
-    if isinstance(camera_matrix, np.ndarray):
-        # Construct transformation matrix to convert from OpenGL space to Blender space
-        flip_yz = np.array([[1, 0, 0, 0], [0, 0, -1, 0], [0, 1, 0, 0], [0, 0, 0, 1]])
-        camera_matrix_blender = np.dot(flip_yz, camera_matrix)
-    else:
-        # Construct transformation matrix to convert from OpenGL space to Blender space
-        flip_yz = torch.tensor(
-            [[1, 0, 0, 0], [0, 0, -1, 0], [0, 1, 0, 0], [0, 0, 0, 1]]
-        )
-        if camera_matrix.ndim == 3:
-            flip_yz = flip_yz.unsqueeze(0)
-        camera_matrix_blender = torch.matmul(flip_yz.to(camera_matrix), camera_matrix)
-    return camera_matrix_blender
-
-
-def get_camera(
-    num_frames, elevation=15, azimuth_start=0, azimuth_span=360, blender_coord=True, extra_view=False,
-):
-    angle_gap = azimuth_span / num_frames
-    cameras = []
-    for azimuth in np.arange(azimuth_start, azimuth_span + azimuth_start, angle_gap):
-        camera_matrix = create_camera_to_world_matrix(elevation, azimuth)
-        if blender_coord:
-            camera_matrix = convert_opengl_to_blender(camera_matrix)
-        cameras.append(camera_matrix.flatten())
-    if extra_view:
-        dim = len(cameras[0])
-        cameras.append(np.zeros(dim))
-    return torch.tensor(np.stack(cameras, 0)).float()
-
-
-class ImageDreamPipeline(DiffusionPipeline):
-    def __init__(
-        self,
-        vae: AutoencoderKL,
-        unet: MultiViewUNetModel,
-        tokenizer: CLIPTokenizer,
-        text_encoder: CLIPTextModel,
-        scheduler: DDIMScheduler,
-        feature_extractor: CLIPFeatureExtractor = None,
-        image_encoder: CLIPVisionModel = None,
-        requires_safety_checker: bool = False,
-    ):
-        super().__init__()
-
-        if hasattr(scheduler.config, "steps_offset") and scheduler.config.steps_offset != 1:  # type: ignore
-            deprecation_message = (
-                f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`"
-                f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure "  # type: ignore
-                "to update the config accordingly as leaving `steps_offset` might led to incorrect results"
-                " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub,"
-                " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`"
-                " file"
-            )
-            deprecate(
-                "steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False
-            )
-            new_config = dict(scheduler.config)
-            new_config["steps_offset"] = 1
-            scheduler._internal_dict = FrozenDict(new_config)
-
-        if hasattr(scheduler.config, "clip_sample") and scheduler.config.clip_sample is True:  # type: ignore
-            deprecation_message = (
-                f"The configuration file of this scheduler: {scheduler} has not set the configuration `clip_sample`."
-                " `clip_sample` should be set to False in the configuration file. Please make sure to update the"
-                " config accordingly as not setting `clip_sample` in the config might lead to incorrect results in"
-                " future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it would be very"
-                " nice if you could open a Pull request for the `scheduler/scheduler_config.json` file"
-            )
-            deprecate(
-                "clip_sample not set", "1.0.0", deprecation_message, standard_warn=False
-            )
-            new_config = dict(scheduler.config)
-            new_config["clip_sample"] = False
-            scheduler._internal_dict = FrozenDict(new_config)
-
-        self.register_modules(
-            vae=vae,
-            unet=unet,
-            scheduler=scheduler,
-            tokenizer=tokenizer,
-            text_encoder=text_encoder,
-            feature_extractor=feature_extractor,
-            image_encoder=image_encoder,
-        )
-        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
-        self.register_to_config(requires_safety_checker=requires_safety_checker)
-
-    def enable_vae_slicing(self):
-        r"""
-        Enable sliced VAE decoding.
-
-        When this option is enabled, the VAE will split the input tensor in slices to compute decoding in several
-        steps. This is useful to save some memory and allow larger batch sizes.
-        """
-        self.vae.enable_slicing()
-
-    def disable_vae_slicing(self):
-        r"""
-        Disable sliced VAE decoding. If `enable_vae_slicing` was previously invoked, this method will go back to
-        computing decoding in one step.
-        """
-        self.vae.disable_slicing()
-
-    def enable_vae_tiling(self):
-        r"""
-        Enable tiled VAE decoding.
-
-        When this option is enabled, the VAE will split the input tensor into tiles to compute decoding and encoding in
-        several steps. This is useful to save a large amount of memory and to allow the processing of larger images.
-        """
-        self.vae.enable_tiling()
-
-    def disable_vae_tiling(self):
-        r"""
-        Disable tiled VAE decoding. If `enable_vae_tiling` was previously invoked, this method will go back to
-        computing decoding in one step.
-        """
-        self.vae.disable_tiling()
-
-    def enable_sequential_cpu_offload(self, gpu_id=0):
-        r"""
-        Offloads all models to CPU using accelerate, significantly reducing memory usage. When called, unet,
-        text_encoder, vae and safety checker have their state dicts saved to CPU and then are moved to a
-        `torch.device('meta') and loaded to GPU only when their specific submodule has its `forward` method called.
-        Note that offloading happens on a submodule basis. Memory savings are higher than with
-        `enable_model_cpu_offload`, but performance is lower.
-        """
-        if is_accelerate_available() and is_accelerate_version(">=", "0.14.0"):
-            from accelerate import cpu_offload
-        else:
-            raise ImportError(
-                "`enable_sequential_cpu_offload` requires `accelerate v0.14.0` or higher"
-            )
-
-        device = torch.device(f"cuda:{gpu_id}")
-
-        if self.device.type != "cpu":
-            self.to("cpu", silence_dtype_warnings=True)
-            torch.cuda.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)
-
-        for cpu_offloaded_model in [self.unet, self.text_encoder, self.vae]:
-            cpu_offload(cpu_offloaded_model, device)
-
-    def enable_model_cpu_offload(self, gpu_id=0):
-        r"""
-        Offloads all models to CPU using accelerate, reducing memory usage with a low impact on performance. Compared
-        to `enable_sequential_cpu_offload`, this method moves one whole model at a time to the GPU when its `forward`
-        method is called, and the model remains in GPU until the next model runs. Memory savings are lower than with
-        `enable_sequential_cpu_offload`, but performance is much better due to the iterative execution of the `unet`.
-        """
-        if is_accelerate_available() and is_accelerate_version(">=", "0.17.0.dev0"):
-            from accelerate import cpu_offload_with_hook
-        else:
-            raise ImportError(
-                "`enable_model_offload` requires `accelerate v0.17.0` or higher."
-            )
-
-        device = torch.device(f"cuda:{gpu_id}")
-
-        if self.device.type != "cpu":
-            self.to("cpu", silence_dtype_warnings=True)
-            torch.cuda.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)
-
-        hook = None
-        for cpu_offloaded_model in [self.text_encoder, self.unet, self.vae]:
-            _, hook = cpu_offload_with_hook(
-                cpu_offloaded_model, device, prev_module_hook=hook
-            )
-
-        # We'll offload the last model manually.
-        self.final_offload_hook = hook
-
-    @property
-    def _execution_device(self):
-        r"""
-        Returns the device on which the pipeline's models will be executed. After calling
-        `pipeline.enable_sequential_cpu_offload()` the execution device can only be inferred from Accelerate's module
-        hooks.
-        """
-        if not hasattr(self.unet, "_hf_hook"):
-            return self.device
-        for module in self.unet.modules():
-            if (
-                hasattr(module, "_hf_hook")
-                and hasattr(module._hf_hook, "execution_device")
-                and module._hf_hook.execution_device is not None
-            ):
-                return torch.device(module._hf_hook.execution_device)
-        return self.device
-
-    def _encode_prompt(
-        self,
-        prompt,
-        device,
-        num_images_per_prompt,
-        do_classifier_free_guidance: bool,
-        negative_prompt=None,
-    ):
-        r"""
-        Encodes the prompt into text encoder hidden states.
-
-        Args:
-             prompt (`str` or `List[str]`, *optional*):
-                prompt to be encoded
-            device: (`torch.device`):
-                torch device
-            num_images_per_prompt (`int`):
-                number of images that should be generated per prompt
-            do_classifier_free_guidance (`bool`):
-                whether to use classifier free guidance or not
-            negative_prompt (`str` or `List[str]`, *optional*):
-                The prompt or prompts not to guide the image generation. If not defined, one has to pass
-                `negative_prompt_embeds`. instead. If not defined, one has to pass `negative_prompt_embeds`. instead.
-                Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`).
-            prompt_embeds (`torch.FloatTensor`, *optional*):
-                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
-                provided, text embeddings will be generated from `prompt` input argument.
-            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
-                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
-                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
-                argument.
-        """
-        if prompt is not None and isinstance(prompt, str):
-            batch_size = 1
-        elif prompt is not None and isinstance(prompt, list):
-            batch_size = len(prompt)
-        else:
-            raise ValueError(
-                f"`prompt` should be either a string or a list of strings, but got {type(prompt)}."
-            )
-
-        text_inputs = self.tokenizer(
-            prompt,
-            padding="max_length",
-            max_length=self.tokenizer.model_max_length,
-            truncation=True,
-            return_tensors="pt",
-        )
-        text_input_ids = text_inputs.input_ids
-        untruncated_ids = self.tokenizer(
-            prompt, padding="longest", return_tensors="pt"
-        ).input_ids
-
-        if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
-            text_input_ids, untruncated_ids
-        ):
-            removed_text = self.tokenizer.batch_decode(
-                untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1]
-            )
-            logger.warning(
-                "The following part of your input was truncated because CLIP can only handle sequences up to"
-                f" {self.tokenizer.model_max_length} tokens: {removed_text}"
-            )
-
-        if (
-            hasattr(self.text_encoder.config, "use_attention_mask")
-            and self.text_encoder.config.use_attention_mask
-        ):
-            attention_mask = text_inputs.attention_mask.to(device)
-        else:
-            attention_mask = None
-
-        prompt_embeds = self.text_encoder(
-            text_input_ids.to(device),
-            attention_mask=attention_mask,
-        )
-        prompt_embeds = prompt_embeds[0]
-
-        prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device)
-
-        bs_embed, seq_len, _ = prompt_embeds.shape
-        # duplicate text embeddings for each generation per prompt, using mps friendly method
-        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
-        prompt_embeds = prompt_embeds.view(
-            bs_embed * num_images_per_prompt, seq_len, -1
-        )
-
-        # get unconditional embeddings for classifier free guidance
-        if do_classifier_free_guidance:
-            uncond_tokens: List[str]
-            if negative_prompt is None:
-                uncond_tokens = [""] * batch_size
-            elif type(prompt) is not type(negative_prompt):
-                raise TypeError(
-                    f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
-                    f" {type(prompt)}."
-                )
-            elif isinstance(negative_prompt, str):
-                uncond_tokens = [negative_prompt]
-            elif batch_size != len(negative_prompt):
-                raise ValueError(
-                    f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
-                    f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
-                    " the batch size of `prompt`."
-                )
-            else:
-                uncond_tokens = negative_prompt
-
-            max_length = prompt_embeds.shape[1]
-            uncond_input = self.tokenizer(
-                uncond_tokens,
-                padding="max_length",
-                max_length=max_length,
-                truncation=True,
-                return_tensors="pt",
-            )
-
-            if (
-                hasattr(self.text_encoder.config, "use_attention_mask")
-                and self.text_encoder.config.use_attention_mask
-            ):
-                attention_mask = uncond_input.attention_mask.to(device)
-            else:
-                attention_mask = None
-
-            negative_prompt_embeds = self.text_encoder(
-                uncond_input.input_ids.to(device),
-                attention_mask=attention_mask,
-            )
-            negative_prompt_embeds = negative_prompt_embeds[0]
-
-            # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
-            seq_len = negative_prompt_embeds.shape[1]
-
-            negative_prompt_embeds = negative_prompt_embeds.to(
-                dtype=self.text_encoder.dtype, device=device
-            )
-
-            negative_prompt_embeds = negative_prompt_embeds.repeat(
-                1, num_images_per_prompt, 1
-            )
-            negative_prompt_embeds = negative_prompt_embeds.view(
-                batch_size * num_images_per_prompt, seq_len, -1
-            )
-
-            # For classifier free guidance, we need to do two forward passes.
-            # Here we concatenate the unconditional and text embeddings into a single batch
-            # to avoid doing two forward passes
-            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
-
-        return prompt_embeds
-
-    def decode_latents(self, latents):
-        latents = 1 / self.vae.config.scaling_factor * latents
-        image = self.vae.decode(latents).sample
-        image = (image / 2 + 0.5).clamp(0, 1)
-        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
-        image = image.cpu().permute(0, 2, 3, 1).float().numpy()
-        return image
-
-    def prepare_extra_step_kwargs(self, generator, eta):
-        # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
-        # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
-        # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
-        # and should be between [0, 1]
-
-        accepts_eta = "eta" in set(
-            inspect.signature(self.scheduler.step).parameters.keys()
-        )
-        extra_step_kwargs = {}
-        if accepts_eta:
-            extra_step_kwargs["eta"] = eta
-
-        # check if the scheduler accepts generator
-        accepts_generator = "generator" in set(
-            inspect.signature(self.scheduler.step).parameters.keys()
-        )
-        if accepts_generator:
-            extra_step_kwargs["generator"] = generator
-        return extra_step_kwargs
-
-    def prepare_latents(
-        self,
-        batch_size,
-        num_channels_latents,
-        height,
-        width,
-        dtype,
-        device,
-        generator,
-        latents=None,
-    ):
-        shape = (
-            batch_size,
-            num_channels_latents,
-            height // self.vae_scale_factor,
-            width // self.vae_scale_factor,
-        )
-        if isinstance(generator, list) and len(generator) != batch_size:
-            raise ValueError(
-                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
-                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
-            )
-
-        if latents is None:
-            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
-
-    def encode_image(self, image, device, num_images_per_prompt):
-        dtype = next(self.image_encoder.parameters()).dtype
-
-        image = (image * 255).astype(np.uint8)
-        image = self.feature_extractor(image, return_tensors="pt").pixel_values
-        
-        image = image.to(device=device, dtype=dtype)
-        
-        image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2]
-        image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0)
-
-        # imagedream directly use zero as uncond image embeddings
-        uncond_image_enc_hidden_states = torch.zeros_like(image_enc_hidden_states)
-
-        return uncond_image_enc_hidden_states, image_enc_hidden_states
-
-    def encode_image_latents(self, image, device, num_images_per_prompt):
-        
-        image = torch.from_numpy(image).to(device)
-        posterior = self.vae.encode(image).latent_dist
-
-        latents = posterior.sample() * self.vae.config.scaling_factor # [B, C, H, W]
-        latents = latents.repeat_interleave(num_images_per_prompt, dim=0)
-
-        return torch.zeros_like(latents), latents
-
-    @torch.no_grad()
-    def __call__(
-        self,
-        image, # input image, np.ndarray float32!
-        prompt: str = "a car",
-        height: int = 256,
-        width: int = 256,
-        num_inference_steps: int = 50,
-        guidance_scale: float = 7.0,
-        negative_prompt: str = "bad quality",
-        num_images_per_prompt: int = 1,
-        eta: float = 0.0,
-        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
-        output_type: Optional[str] = "image",
-        callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
-        callback_steps: int = 1,
-        num_frames: int = 4,
-        device=torch.device("cuda:0"),
-    ):
-        self.unet = self.unet.to(device=device)
-        self.vae = self.vae.to(device=device)
-
-        self.text_encoder = self.text_encoder.to(device=device)
-
-        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
-        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
-        # corresponds to doing no classifier free guidance.
-        do_classifier_free_guidance = guidance_scale > 1.0
-
-        # Prepare timesteps
-        self.scheduler.set_timesteps(num_inference_steps, device=device)
-        timesteps = self.scheduler.timesteps
-
-        # encode image
-        assert isinstance(image, np.ndarray) and image.dtype == np.float32
-
-        self.image_encoder = self.image_encoder.to(device=device)
-        image_embeds_neg, image_embeds_pos = self.encode_image(image, device, num_images_per_prompt)
-        kiui.lo(image_embeds_pos) # should be [1, 257, 1280]?
-
-        image_latents_neg, image_latents_pos = self.encode_image_latents(image, device, num_images_per_prompt)
-        kiui.lo(image_latents_pos)
-        
-        # encode text
-        _prompt_embeds = self._encode_prompt(
-            prompt=prompt,
-            device=device,
-            num_images_per_prompt=num_images_per_prompt,
-            do_classifier_free_guidance=do_classifier_free_guidance,
-            negative_prompt=negative_prompt,
-        )  # type: ignore
-        prompt_embeds_neg, prompt_embeds_pos = _prompt_embeds.chunk(2)
-
-        # Prepare latent variables
-        latents: torch.Tensor = self.prepare_latents(
-            (num_frames + 1) * num_images_per_prompt,
-            4, # channel
-            height,
-            width,
-            prompt_embeds_pos.dtype,
-            device,
-            generator,
-            None,
-        )
-
-        camera = get_camera(num_frames, extra_view=True).to(dtype=latents.dtype, device=device)
-        camera = camera.repeat(num_images_per_prompt, 1).to(self.device)
-
-        # Prepare extra step kwargs.
-        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
-
-        # Denoising loop
-        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
-        with self.progress_bar(total=num_inference_steps) as progress_bar:
-            for i, t in enumerate(timesteps):
-                # expand the latents if we are doing classifier free guidance
-                multiplier = 2 if do_classifier_free_guidance else 1
-                latent_model_input = torch.cat([latents] * multiplier)
-                latent_model_input = self.scheduler.scale_model_input(
-                    latent_model_input, t
-                )
-
-                # predict the noise residual
-                noise_pred = self.unet.forward(
-                    x=latent_model_input,
-                    timesteps=torch.tensor(
-                        [t] * (num_frames + 1) * multiplier,
-                        dtype=latent_model_input.dtype,
-                        device=device,
-                    ),
-                    context=torch.cat(
-                        [prompt_embeds_neg] * (num_frames + 1) + [prompt_embeds_pos] * (num_frames + 1)
-                    ),
-                    num_frames=num_frames + 1,
-                    camera=torch.cat([camera] * multiplier),
-                    # for with_ip
-                    ip=torch.cat(
-                        [image_embeds_neg] * (num_frames + 1) + [image_embeds_pos] * (num_frames + 1)
-                    ),
-                    ip_img=torch.cat(
-                        [image_latents_neg] * (num_frames + 1) + [image_latents_pos] * (num_frames + 1)
-                    ),
-                )
-
-                # perform guidance
-                if do_classifier_free_guidance:
-                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
-                    noise_pred = noise_pred_uncond + guidance_scale * (
-                        noise_pred_text - noise_pred_uncond
-                    )
-
-                # compute the previous noisy sample x_t -> x_t-1
-                latents: torch.Tensor = self.scheduler.step(
-                    noise_pred, t, latents, **extra_step_kwargs, return_dict=False
-                )[0]
-
-                # call the callback, if provided
-                if i == len(timesteps) - 1 or (
-                    (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0
-                ):
-                    progress_bar.update()
-                    if callback is not None and i % callback_steps == 0:
-                        callback(i, t, latents)  # type: ignore
-
-        # Post-processing
-        if output_type == "latent":
-            image = latents
-        elif output_type == "pil":
-            image = self.decode_latents(latents)
-            image = self.numpy_to_pil(image)
-        else:
-            image = self.decode_latents(latents)
-
-        # Offload last model to CPU
-        if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
-            self.final_offload_hook.offload()
-
-        return image

imagedream/util.py

@@ -1,116 +0,0 @@
-import math
-import torch
-import torch.nn as nn
-from einops import repeat
-
-
-def checkpoint(func, inputs, params, flag):
-    """
-    Evaluate a function without caching intermediate activations, allowing for
-    reduced memory at the expense of extra compute in the backward pass.
-    :param func: the function to evaluate.
-    :param inputs: the argument sequence to pass to `func`.
-    :param params: a sequence of parameters `func` depends on but does not
-                   explicitly take as arguments.
-    :param flag: if False, disable gradient checkpointing.
-    """
-    if flag:
-        args = tuple(inputs) + tuple(params)
-        return CheckpointFunction.apply(func, len(inputs), *args)
-    else:
-        return func(*inputs)
-
-
-class CheckpointFunction(torch.autograd.Function):
-    @staticmethod
-    def forward(ctx, run_function, length, *args):
-        ctx.run_function = run_function
-        ctx.input_tensors = list(args[:length])
-        ctx.input_params = list(args[length:])
-
-        with torch.no_grad():
-            output_tensors = ctx.run_function(*ctx.input_tensors)
-        return output_tensors
-
-    @staticmethod
-    def backward(ctx, *output_grads):
-        ctx.input_tensors = [x.detach().requires_grad_(True) for x in ctx.input_tensors]
-        with torch.enable_grad():
-            # Fixes a bug where the first op in run_function modifies the
-            # Tensor storage in place, which is not allowed for detach()'d
-            # Tensors.
-            shallow_copies = [x.view_as(x) for x in ctx.input_tensors]
-            output_tensors = ctx.run_function(*shallow_copies)
-        input_grads = torch.autograd.grad(
-            output_tensors,
-            ctx.input_tensors + ctx.input_params,
-            output_grads,
-            allow_unused=True,
-        )
-        del ctx.input_tensors
-        del ctx.input_params
-        del output_tensors
-        return (None, None) + input_grads
-
-
-def timestep_embedding(timesteps, dim, max_period=10000, repeat_only=False):
-    """
-    Create sinusoidal timestep embeddings.
-    :param timesteps: a 1-D Tensor of N indices, one per batch element.
-                      These may be fractional.
-    :param dim: the dimension of the output.
-    :param max_period: controls the minimum frequency of the embeddings.
-    :return: an [N x dim] Tensor of positional embeddings.
-    """
-    if not repeat_only:
-        half = dim // 2
-        freqs = torch.exp(
-            -math.log(max_period)
-            * torch.arange(start=0, end=half, dtype=torch.float32)
-            / half
-        ).to(device=timesteps.device)
-        args = timesteps[:, None] * freqs[None]
-        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
-        if dim % 2:
-            embedding = torch.cat(
-                [embedding, torch.zeros_like(embedding[:, :1])], dim=-1
-            )
-    else:
-        embedding = repeat(timesteps, "b -> b d", d=dim)
-    # import pdb; pdb.set_trace()
-    return embedding
-
-
-def zero_module(module):
-    """
-    Zero out the parameters of a module and return it.
-    """
-    for p in module.parameters():
-        p.detach().zero_()
-    return module
-
-
-def conv_nd(dims, *args, **kwargs):
-    """
-    Create a 1D, 2D, or 3D convolution module.
-    """
-    if dims == 1:
-        return nn.Conv1d(*args, **kwargs)
-    elif dims == 2:
-        return nn.Conv2d(*args, **kwargs)
-    elif dims == 3:
-        return nn.Conv3d(*args, **kwargs)
-    raise ValueError(f"unsupported dimensions: {dims}")
-
-
-def avg_pool_nd(dims, *args, **kwargs):
-    """
-    Create a 1D, 2D, or 3D average pooling module.
-    """
-    if dims == 1:
-        return nn.AvgPool1d(*args, **kwargs)
-    elif dims == 2:
-        return nn.AvgPool2d(*args, **kwargs)
-    elif dims == 3:
-        return nn.AvgPool3d(*args, **kwargs)
-    raise ValueError(f"unsupported dimensions: {dims}")

mvdream/attention.py

@@ -1,26 +1,16 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from torch.amp.autocast_mode import autocast
 
 from inspect import isfunction
 from einops import rearrange, repeat
 from typing import Optional, Any
-from .util import checkpoint, zero_module
-
-try:
-    import xformers  # type: ignore
-    import xformers.ops  # type: ignore
-    XFORMERS_IS_AVAILBLE = True
-except:
-    print(f'[WARN] xformers is unavailable!')
-    XFORMERS_IS_AVAILBLE = False
 
-# CrossAttn precision handling
-import os
-
-_ATTN_PRECISION = os.environ.get("ATTN_PRECISION", "fp32")
+# require xformers
+import xformers  # type: ignore
+import xformers.ops  # type: ignore
 
+from .util import checkpoint, zero_module
 
 def default(val, d):
     if val is not None:
@@ -57,68 +47,33 @@ class FeedForward(nn.Module):
         return self.net(x)
 
 
-class CrossAttention(nn.Module):
-    def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0.0):
-        super().__init__()
-        inner_dim = dim_head * heads
-        context_dim = default(context_dim, query_dim)
-
-        self.scale = dim_head**-0.5
-        self.heads = heads
-
-        self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
-        self.to_k = nn.Linear(context_dim, inner_dim, bias=False)
-        self.to_v = nn.Linear(context_dim, inner_dim, bias=False)
-
-        self.to_out = nn.Sequential(
-            nn.Linear(inner_dim, query_dim), nn.Dropout(dropout)
-        )
-
-    def forward(self, x, context=None, mask=None):
-        h = self.heads
-
-        q = self.to_q(x)
-        context = default(context, x)
-        k = self.to_k(context)
-        v = self.to_v(context)
-
-        q, k, v = map(lambda t: rearrange(t, "b n (h d) -> (b h) n d", h=h), (q, k, v))
-
-        # force cast to fp32 to avoid overflowing
-        if _ATTN_PRECISION == "fp32":
-            with autocast(enabled=False, device_type="cuda"):
-                q, k = q.float(), k.float()
-                sim = torch.einsum("b i d, b j d -> b i j", q, k) * self.scale
-        else:
-            sim = torch.einsum("b i d, b j d -> b i j", q, k) * self.scale
-
-        del q, k
-
-        if mask is not None:
-            mask = rearrange(mask, "b ... -> b (...)")
-            max_neg_value = -torch.finfo(sim.dtype).max
-            mask = repeat(mask, "b j -> (b h) () j", h=h)
-            sim.masked_fill_(~mask, max_neg_value)
-
-        # attention, what we cannot get enough of
-        sim = sim.softmax(dim=-1)
-
-        out = torch.einsum("b i j, b j d -> b i d", sim, v)
-        out = rearrange(out, "(b h) n d -> b n (h d)", h=h)
-        return self.to_out(out)
-
-
 class MemoryEfficientCrossAttention(nn.Module):
     # https://github.com/MatthieuTPHR/diffusers/blob/d80b531ff8060ec1ea982b65a1b8df70f73aa67c/src/diffusers/models/attention.py#L223
-    def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0.0):
+    def __init__(
+            self, 
+            query_dim, 
+            context_dim=None, 
+            heads=8, 
+            dim_head=64, 
+            dropout=0.0,
+            ip_dim=0,
+            ip_weight=1,
+        ):
         super().__init__()
-        # print(f"Setting up {self.__class__.__name__}. Query dim is {query_dim}, context_dim is {context_dim} and using {heads} heads.")
+        
         inner_dim = dim_head * heads
         context_dim = default(context_dim, query_dim)
 
         self.heads = heads
         self.dim_head = dim_head
 
+        self.ip_dim = ip_dim
+        self.ip_weight = ip_weight
+
+        if self.ip_dim > 0:
+            self.to_k_ip = nn.Linear(context_dim, inner_dim, bias=False)
+            self.to_v_ip = nn.Linear(context_dim, inner_dim, bias=False)
+
         self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
         self.to_k = nn.Linear(context_dim, inner_dim, bias=False)
         self.to_v = nn.Linear(context_dim, inner_dim, bias=False)
@@ -128,9 +83,18 @@ class MemoryEfficientCrossAttention(nn.Module):
         )
         self.attention_op: Optional[Any] = None
 
-    def forward(self, x, context=None, mask=None):
+    def forward(self, x, context=None):
         q = self.to_q(x)
         context = default(context, x)
+
+        if self.ip_dim > 0:
+            # context dim [(b frame_num), (77 + img_token), 1024]
+            token_len = context.shape[1]
+            context_ip = context[:, -self.ip_dim :, :]
+            k_ip = self.to_k_ip(context_ip)
+            v_ip = self.to_v_ip(context_ip)
+            context = context[:, : (token_len - self.ip_dim), :]
+
         k = self.to_k(context)
         v = self.to_v(context)
 
@@ -149,8 +113,21 @@ class MemoryEfficientCrossAttention(nn.Module):
             q, k, v, attn_bias=None, op=self.attention_op
         )
 
-        if mask is not None:
-            raise NotImplementedError
+        if self.ip_dim > 0:
+            k_ip, v_ip = map(
+                lambda t: t.unsqueeze(3)
+                .reshape(b, t.shape[1], self.heads, self.dim_head)
+                .permute(0, 2, 1, 3)
+                .reshape(b * self.heads, t.shape[1], self.dim_head)
+                .contiguous(),
+                (k_ip, v_ip),
+            )
+            # actually compute the attention, what we cannot get enough of
+            out_ip = xformers.ops.memory_efficient_attention(
+                q, k_ip, v_ip, attn_bias=None, op=self.attention_op
+            )
+            out = out + self.ip_weight * out_ip
+
         out = (
             out.unsqueeze(0)
             .reshape(b, self.heads, out.shape[1], self.dim_head)
@@ -160,148 +137,45 @@ class MemoryEfficientCrossAttention(nn.Module):
         return self.to_out(out)
 
 
-class BasicTransformerBlock(nn.Module):
-    ATTENTION_MODES = {
-        "softmax": CrossAttention,
-        "softmax-xformers": MemoryEfficientCrossAttention,
-    }  # vanilla attention
-
+class BasicTransformerBlock3D(nn.Module):
+    
     def __init__(
         self,
         dim,
         n_heads,
         d_head,
+        context_dim,
         dropout=0.0,
-        context_dim=None,
         gated_ff=True,
         checkpoint=True,
-        disable_self_attn=False,
+        ip_dim=0,
+        ip_weight=1,
     ):
         super().__init__()
-        attn_mode = "softmax-xformers" if XFORMERS_IS_AVAILBLE else "softmax"
-        assert attn_mode in self.ATTENTION_MODES
-        attn_cls = self.ATTENTION_MODES[attn_mode]
-        self.disable_self_attn = disable_self_attn
-        self.attn1 = attn_cls(
+
+        self.attn1 = MemoryEfficientCrossAttention(
             query_dim=dim,
+            context_dim=None, # self-attention
             heads=n_heads,
             dim_head=d_head,
             dropout=dropout,
-            context_dim=context_dim if self.disable_self_attn else None,
-        )  # is a self-attention if not self.disable_self_attn
+        )
         self.ff = FeedForward(dim, dropout=dropout, glu=gated_ff)
-        self.attn2 = attn_cls(
+        self.attn2 = MemoryEfficientCrossAttention(
             query_dim=dim,
             context_dim=context_dim,
             heads=n_heads,
             dim_head=d_head,
             dropout=dropout,
-        )  # is self-attn if context is none
+            # ip only applies to cross-attention
+            ip_dim=ip_dim,
+            ip_weight=ip_weight,
+        ) 
         self.norm1 = nn.LayerNorm(dim)
         self.norm2 = nn.LayerNorm(dim)
         self.norm3 = nn.LayerNorm(dim)
         self.checkpoint = checkpoint
 
-    def forward(self, x, context=None):
-        return checkpoint(
-            self._forward, (x, context), self.parameters(), self.checkpoint
-        )
-
-    def _forward(self, x, context=None):
-        x = (
-            self.attn1(
-                self.norm1(x), context=context if self.disable_self_attn else None
-            )
-            + x
-        )
-        x = self.attn2(self.norm2(x), context=context) + x
-        x = self.ff(self.norm3(x)) + x
-        return x
-
-
-class SpatialTransformer(nn.Module):
-    """
-    Transformer block for image-like data.
-    First, project the input (aka embedding)
-    and reshape to b, t, d.
-    Then apply standard transformer action.
-    Finally, reshape to image
-    NEW: use_linear for more efficiency instead of the 1x1 convs
-    """
-
-    def __init__(
-        self,
-        in_channels,
-        n_heads,
-        d_head,
-        depth=1,
-        dropout=0.0,
-        context_dim=None,
-        disable_self_attn=False,
-        use_linear=False,
-        use_checkpoint=True,
-    ):
-        super().__init__()
-        assert context_dim is not None
-        if not isinstance(context_dim, list):
-            context_dim = [context_dim]
-        self.in_channels = in_channels
-        inner_dim = n_heads * d_head
-        self.norm = nn.GroupNorm(
-            num_groups=32, num_channels=in_channels, eps=1e-6, affine=True
-        )
-        if not use_linear:
-            self.proj_in = nn.Conv2d(
-                in_channels, inner_dim, kernel_size=1, stride=1, padding=0
-            )
-        else:
-            self.proj_in = nn.Linear(in_channels, inner_dim)
-
-        self.transformer_blocks = nn.ModuleList(
-            [
-                BasicTransformerBlock(
-                    inner_dim,
-                    n_heads,
-                    d_head,
-                    dropout=dropout,
-                    context_dim=context_dim[d],
-                    disable_self_attn=disable_self_attn,
-                    checkpoint=use_checkpoint,
-                )
-                for d in range(depth)
-            ]
-        )
-        if not use_linear:
-            self.proj_out = zero_module(
-                nn.Conv2d(inner_dim, in_channels, kernel_size=1, stride=1, padding=0)
-            )
-        else:
-            self.proj_out = zero_module(nn.Linear(in_channels, inner_dim))
-        self.use_linear = use_linear
-
-    def forward(self, x, context=None):
-        # note: if no context is given, cross-attention defaults to self-attention
-        if not isinstance(context, list):
-            context = [context]
-        b, c, h, w = x.shape
-        x_in = x
-        x = self.norm(x)
-        if not self.use_linear:
-            x = self.proj_in(x)
-        x = rearrange(x, "b c h w -> b (h w) c").contiguous()
-        if self.use_linear:
-            x = self.proj_in(x)
-        for i, block in enumerate(self.transformer_blocks):
-            x = block(x, context=context[i])
-        if self.use_linear:
-            x = self.proj_out(x)
-        x = rearrange(x, "b (h w) c -> b c h w", h=h, w=w).contiguous()
-        if not self.use_linear:
-            x = self.proj_out(x)
-        return x + x_in
-
-
-class BasicTransformerBlock3D(BasicTransformerBlock):
     def forward(self, x, context=None, num_frames=1):
         return checkpoint(
             self._forward, (x, context, num_frames), self.parameters(), self.checkpoint
@@ -309,12 +183,7 @@ class BasicTransformerBlock3D(BasicTransformerBlock):
 
     def _forward(self, x, context=None, num_frames=1):
         x = rearrange(x, "(b f) l c -> b (f l) c", f=num_frames).contiguous()
-        x = (
-            self.attn1(
-                self.norm1(x), context=context if self.disable_self_attn else None
-            )
-            + x
-        )
+        x = self.attn1(self.norm1(x), context=None) + x
         x = rearrange(x, "b (f l) c -> (b f) l c", f=num_frames).contiguous()
         x = self.attn2(self.norm2(x), context=context) + x
         x = self.ff(self.norm3(x)) + x
@@ -322,35 +191,31 @@ class BasicTransformerBlock3D(BasicTransformerBlock):
 
 
 class SpatialTransformer3D(nn.Module):
-    """3D self-attention"""
 
     def __init__(
         self,
         in_channels,
         n_heads,
         d_head,
+        context_dim, # cross attention input dim
         depth=1,
         dropout=0.0,
-        context_dim=None,
-        disable_self_attn=False,
-        use_linear=True,
+        ip_dim=0,
+        ip_weight=1,
         use_checkpoint=True,
     ):
         super().__init__()
-        assert context_dim is not None
+
         if not isinstance(context_dim, list):
             context_dim = [context_dim]
+
         self.in_channels = in_channels
+
         inner_dim = n_heads * d_head
         self.norm = nn.GroupNorm(
             num_groups=32, num_channels=in_channels, eps=1e-6, affine=True
         )
-        if not use_linear:
-            self.proj_in = nn.Conv2d(
-                in_channels, inner_dim, kernel_size=1, stride=1, padding=0
-            )
-        else:
-            self.proj_in = nn.Linear(in_channels, inner_dim)
+        self.proj_in = nn.Linear(in_channels, inner_dim)
 
         self.transformer_blocks = nn.ModuleList(
             [
@@ -358,21 +223,18 @@ class SpatialTransformer3D(nn.Module):
                     inner_dim,
                     n_heads,
                     d_head,
-                    dropout=dropout,
                     context_dim=context_dim[d],
-                    disable_self_attn=disable_self_attn,
+                    dropout=dropout,
                     checkpoint=use_checkpoint,
+                    ip_dim=ip_dim,
+                    ip_weight=ip_weight,
                 )
                 for d in range(depth)
             ]
         )
-        if not use_linear:
-            self.proj_out = zero_module(
-                nn.Conv2d(inner_dim, in_channels, kernel_size=1, stride=1, padding=0)
-            )
-        else:
-            self.proj_out = zero_module(nn.Linear(in_channels, inner_dim))
-        self.use_linear = use_linear
+        
+        self.proj_out = zero_module(nn.Linear(in_channels, inner_dim))
+        
 
     def forward(self, x, context=None, num_frames=1):
         # note: if no context is given, cross-attention defaults to self-attention
@@ -381,16 +243,11 @@ class SpatialTransformer3D(nn.Module):
         b, c, h, w = x.shape
         x_in = x
         x = self.norm(x)
-        if not self.use_linear:
-            x = self.proj_in(x)
         x = rearrange(x, "b c h w -> b (h w) c").contiguous()
-        if self.use_linear:
-            x = self.proj_in(x)
+        x = self.proj_in(x)
         for i, block in enumerate(self.transformer_blocks):
             x = block(x, context=context[i], num_frames=num_frames)
-        if self.use_linear:
-            x = self.proj_out(x)
+        x = self.proj_out(x)
         x = rearrange(x, "b (h w) c -> b c h w", h=h, w=w).contiguous()
-        if not self.use_linear:
-            x = self.proj_out(x)
+        
         return x + x_in

mvdream/models.py

@@ -13,8 +13,10 @@ from .util import (
     zero_module,
     timestep_embedding,
 )
-from .attention import SpatialTransformer, SpatialTransformer3D
+from .attention import SpatialTransformer3D
+from .adaptor import Resampler, ImageProjModel
 
+import kiui
 
 class CondSequential(nn.Sequential):
     """
@@ -28,8 +30,6 @@ class CondSequential(nn.Sequential):
                 x = layer(x, emb)
             elif isinstance(layer, SpatialTransformer3D):
                 x = layer(x, context, num_frames=num_frames)
-            elif isinstance(layer, SpatialTransformer):
-                x = layer(x, context)
             else:
                 x = layer(x)
         return x
@@ -274,6 +274,8 @@ class MultiViewUNetModel(ModelMixin, ConfigMixin):
         disable_middle_self_attn=False,
         adm_in_channels=None,
         camera_dim=None,
+        ip_dim=0,
+        ip_weight=1.0,
         **kwargs,
     ):
         super().__init__()
@@ -305,9 +307,7 @@ class MultiViewUNetModel(ModelMixin, ConfigMixin):
                     "as a list/tuple (per-level) with the same length as channel_mult"
                 )
             self.num_res_blocks = num_res_blocks
-        if disable_self_attentions is not None:
-            # should be a list of booleans, indicating whether to disable self-attention in TransformerBlocks or not
-            assert len(disable_self_attentions) == len(channel_mult)
+        
         if num_attention_blocks is not None:
             assert len(num_attention_blocks) == len(self.num_res_blocks)
             assert all(
@@ -334,6 +334,21 @@ class MultiViewUNetModel(ModelMixin, ConfigMixin):
         self.num_heads_upsample = num_heads_upsample
         self.predict_codebook_ids = n_embed is not None
 
+        self.ip_dim = ip_dim
+        self.ip_weight = ip_weight
+
+        if self.ip_dim > 0:
+            self.image_embed = Resampler(
+                dim=context_dim,
+                depth=4,
+                dim_head=64,
+                heads=12,
+                num_queries=ip_dim,  # num token
+                embedding_dim=1280,
+                output_dim=context_dim,
+                ff_mult=4,
+            )
+
         time_embed_dim = model_channels * 4
         self.time_embed = nn.Sequential(
             nn.Linear(model_channels, time_embed_dim),
@@ -398,11 +413,6 @@ class MultiViewUNetModel(ModelMixin, ConfigMixin):
                     else:
                         num_heads = ch // num_head_channels
                         dim_head = num_head_channels
-                    
-                    if disable_self_attentions is not None:
-                        disabled_sa = disable_self_attentions[level]
-                    else:
-                        disabled_sa = False
 
                     if num_attention_blocks is None or nr < num_attention_blocks[level]:
                         layers.append(
@@ -410,10 +420,11 @@ class MultiViewUNetModel(ModelMixin, ConfigMixin):
                                 ch,
                                 num_heads,
                                 dim_head,
-                                depth=transformer_depth,
                                 context_dim=context_dim,
-                                disable_self_attn=disabled_sa,
+                                depth=transformer_depth,
                                 use_checkpoint=use_checkpoint,
+                                ip_dim=self.ip_dim,
+                                ip_weight=self.ip_weight,
                             )
                         )
                 self.input_blocks.append(CondSequential(*layers))
@@ -463,10 +474,11 @@ class MultiViewUNetModel(ModelMixin, ConfigMixin):
                 ch,
                 num_heads,
                 dim_head,
-                depth=transformer_depth,
                 context_dim=context_dim,
-                disable_self_attn=disable_middle_self_attn,
+                depth=transformer_depth,
                 use_checkpoint=use_checkpoint,
+                ip_dim=self.ip_dim,
+                ip_weight=self.ip_weight,
             ), 
             ResBlock(
                 ch,
@@ -501,11 +513,6 @@ class MultiViewUNetModel(ModelMixin, ConfigMixin):
                     else:
                         num_heads = ch // num_head_channels
                         dim_head = num_head_channels
-                    
-                    if disable_self_attentions is not None:
-                        disabled_sa = disable_self_attentions[level]
-                    else:
-                        disabled_sa = False
 
                     if num_attention_blocks is None or i < num_attention_blocks[level]:
                         layers.append(
@@ -513,10 +520,11 @@ class MultiViewUNetModel(ModelMixin, ConfigMixin):
                                 ch,
                                 num_heads,
                                 dim_head,
-                                depth=transformer_depth,
                                 context_dim=context_dim,
-                                disable_self_attn=disabled_sa,
+                                depth=transformer_depth,
                                 use_checkpoint=use_checkpoint,
+                                ip_dim=self.ip_dim,
+                                ip_weight=self.ip_weight,
                             )
                         )
                 if level and i == self.num_res_blocks[level]:
@@ -556,9 +564,11 @@ class MultiViewUNetModel(ModelMixin, ConfigMixin):
         x,
         timesteps=None,
         context=None,
-        y: Optional[Tensor] = None,
+        y=None,
         camera=None,
         num_frames=1,
+        ip=None,
+        ip_img=None,
         **kwargs,
     ):
         """
@@ -572,14 +582,14 @@ class MultiViewUNetModel(ModelMixin, ConfigMixin):
         """
         assert (
             x.shape[0] % num_frames == 0
-        ), "[UNet] input batch size must be dividable by num_frames!"
+        ), "input batch size must be dividable by num_frames!"
         assert (y is not None) == (
             self.num_classes is not None
         ), "must specify y if and only if the model is class-conditional"
+
         hs = []
-        t_emb = timestep_embedding(
-            timesteps, self.model_channels, repeat_only=False
-        ).to(x.dtype)
+
+        t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False).to(x.dtype)
 
         emb = self.time_embed(t_emb)
 
@@ -590,8 +600,13 @@ class MultiViewUNetModel(ModelMixin, ConfigMixin):
 
         # Add camera embeddings
         if camera is not None:
-            assert camera.shape[0] == emb.shape[0]
             emb = emb + self.camera_embed(camera)
+        
+        # imagedream variant
+        if self.ip_dim > 0:
+            x[(num_frames - 1) :: num_frames, :, :, :] = ip_img
+            ip_emb = self.image_embed(ip)
+            context = torch.cat((context, ip_emb), 1)
 
         h = x
         for module in self.input_blocks:

mvdream/pipeline_mvdream.py

@@ -1,8 +1,9 @@
 import torch
+import torch.nn.functional as F
 import inspect
 import numpy as np
 from typing import Callable, List, Optional, Union
-from transformers import CLIPTextModel, CLIPTokenizer
+from transformers import CLIPTextModel, CLIPTokenizer, CLIPVisionModel, CLIPImageProcessor
 from diffusers import AutoencoderKL, DiffusionPipeline
 from diffusers.utils import (
     deprecate,
@@ -15,66 +16,17 @@ from diffusers.schedulers import DDIMScheduler
 from diffusers.utils.torch_utils import randn_tensor
 
 from .models import MultiViewUNetModel
+from .util import get_camera
 
-logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
-
+import kiui
 
-def create_camera_to_world_matrix(elevation, azimuth):
-    elevation = np.radians(elevation)
-    azimuth = np.radians(azimuth)
-    # Convert elevation and azimuth angles to Cartesian coordinates on a unit sphere
-    x = np.cos(elevation) * np.sin(azimuth)
-    y = np.sin(elevation)
-    z = np.cos(elevation) * np.cos(azimuth)
-
-    # Calculate camera position, target, and up vectors
-    camera_pos = np.array([x, y, z])
-    target = np.array([0, 0, 0])
-    up = np.array([0, 1, 0])
-
-    # Construct view matrix
-    forward = target - camera_pos
-    forward /= np.linalg.norm(forward)
-    right = np.cross(forward, up)
-    right /= np.linalg.norm(right)
-    new_up = np.cross(right, forward)
-    new_up /= np.linalg.norm(new_up)
-    cam2world = np.eye(4)
-    cam2world[:3, :3] = np.array([right, new_up, -forward]).T
-    cam2world[:3, 3] = camera_pos
-    return cam2world
-
-
-def convert_opengl_to_blender(camera_matrix):
-    if isinstance(camera_matrix, np.ndarray):
-        # Construct transformation matrix to convert from OpenGL space to Blender space
-        flip_yz = np.array([[1, 0, 0, 0], [0, 0, -1, 0], [0, 1, 0, 0], [0, 0, 0, 1]])
-        camera_matrix_blender = np.dot(flip_yz, camera_matrix)
-    else:
-        # Construct transformation matrix to convert from OpenGL space to Blender space
-        flip_yz = torch.tensor(
-            [[1, 0, 0, 0], [0, 0, -1, 0], [0, 1, 0, 0], [0, 0, 0, 1]]
-        )
-        if camera_matrix.ndim == 3:
-            flip_yz = flip_yz.unsqueeze(0)
-        camera_matrix_blender = torch.matmul(flip_yz.to(camera_matrix), camera_matrix)
-    return camera_matrix_blender
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
 
 
-def get_camera(
-    num_frames, elevation=15, azimuth_start=0, azimuth_span=360, blender_coord=True
-):
-    angle_gap = azimuth_span / num_frames
-    cameras = []
-    for azimuth in np.arange(azimuth_start, azimuth_span + azimuth_start, angle_gap):
-        camera_matrix = create_camera_to_world_matrix(elevation, azimuth)
-        if blender_coord:
-            camera_matrix = convert_opengl_to_blender(camera_matrix)
-        cameras.append(camera_matrix.flatten())
-    return torch.tensor(np.stack(cameras, 0)).float()
+class MVDreamPipeline(DiffusionPipeline):
 
+    _optional_components = ["feature_extractor", "image_encoder"]
 
-class MVDreamPipeline(DiffusionPipeline):
     def __init__(
         self,
         vae: AutoencoderKL,
@@ -82,6 +34,9 @@ class MVDreamPipeline(DiffusionPipeline):
         tokenizer: CLIPTokenizer,
         text_encoder: CLIPTextModel,
         scheduler: DDIMScheduler,
+        # imagedream variant
+        feature_extractor: CLIPImageProcessor,
+        image_encoder: CLIPVisionModel,
         requires_safety_checker: bool = False,
     ):
         super().__init__()
@@ -123,6 +78,8 @@ class MVDreamPipeline(DiffusionPipeline):
             scheduler=scheduler,
             tokenizer=tokenizer,
             text_encoder=text_encoder,
+            feature_extractor=feature_extractor,
+            image_encoder=image_encoder,
         )
         self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
         self.register_to_config(requires_safety_checker=requires_safety_checker)
@@ -445,10 +402,42 @@ class MVDreamPipeline(DiffusionPipeline):
         latents = latents * self.scheduler.init_noise_sigma
         return latents
 
+    def encode_image(self, image, device, num_images_per_prompt):
+        dtype = next(self.image_encoder.parameters()).dtype
+
+        image = (image * 255).astype(np.uint8)
+        image = self.feature_extractor(image, return_tensors="pt").pixel_values
+        
+        image = image.to(device=device, dtype=dtype)
+        
+        image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2]
+        image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0)
+
+        # imagedream directly use zero as uncond image embeddings
+        uncond_image_enc_hidden_states = torch.zeros_like(image_enc_hidden_states)
+
+        return uncond_image_enc_hidden_states, image_enc_hidden_states
+
+    def encode_image_latents(self, image, device, num_images_per_prompt):
+        
+        image = torch.from_numpy(image).unsqueeze(0).permute(0, 3, 1, 2) # [1, 3, H, W]
+        image = image.to(device=device)
+        image = F.interpolate(image, (256, 256), mode='bilinear', align_corners=False)
+        dtype = next(self.image_encoder.parameters()).dtype
+        image = image.to(dtype=dtype)
+
+        posterior = self.vae.encode(image).latent_dist
+
+        latents = posterior.sample() * self.vae.config.scaling_factor # [B, C, H, W]
+        latents = latents.repeat_interleave(num_images_per_prompt, dim=0)
+
+        return torch.zeros_like(latents), latents
+
     @torch.no_grad()
     def __call__(
         self,
         prompt: str = "a car",
+        image: Optional[np.ndarray] = None,
         height: int = 256,
         width: int = 256,
         num_inference_steps: int = 50,
@@ -457,10 +446,10 @@ class MVDreamPipeline(DiffusionPipeline):
         num_images_per_prompt: int = 1,
         eta: float = 0.0,
         generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
-        output_type: Optional[str] = "image",
+        output_type: Optional[str] = "numpy", # pil, numpy, latents
         callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
         callback_steps: int = 1,
-        batch_size: int = 4,
+        num_frames: int = 4,
         device=torch.device("cuda:0"),
     ):
         self.unet = self.unet.to(device=device)
@@ -477,7 +466,15 @@ class MVDreamPipeline(DiffusionPipeline):
         self.scheduler.set_timesteps(num_inference_steps, device=device)
         timesteps = self.scheduler.timesteps
 
-        _prompt_embeds: torch.Tensor = self._encode_prompt(
+        # imagedream variant (TODO: debug)
+        if image is not None:
+            assert isinstance(image, np.ndarray) and image.dtype == np.float32
+
+            self.image_encoder = self.image_encoder.to(device=device)
+            image_embeds_neg, image_embeds_pos = self.encode_image(image, device, num_images_per_prompt)
+            image_latents_neg, image_latents_pos = self.encode_image_latents(image, device, num_images_per_prompt)
+            
+        _prompt_embeds = self._encode_prompt(
             prompt=prompt,
             device=device,
             num_images_per_prompt=num_images_per_prompt,
@@ -487,8 +484,9 @@ class MVDreamPipeline(DiffusionPipeline):
         prompt_embeds_neg, prompt_embeds_pos = _prompt_embeds.chunk(2)
 
         # Prepare latent variables
+        actual_num_frames = num_frames if image is None else num_frames + 1
         latents: torch.Tensor = self.prepare_latents(
-            batch_size * num_images_per_prompt,
+            actual_num_frames * num_images_per_prompt,
             4,
             height,
             width,
@@ -498,9 +496,9 @@ class MVDreamPipeline(DiffusionPipeline):
             None,
         )
 
-        camera = get_camera(batch_size).to(dtype=latents.dtype, device=device)
+        camera = get_camera(num_frames, extra_view=(actual_num_frames != num_frames)).to(dtype=latents.dtype, device=device)
 
-        # Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
+        # Prepare extra step kwargs.
         extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
 
         # Denoising loop
@@ -514,20 +512,21 @@ class MVDreamPipeline(DiffusionPipeline):
                     latent_model_input, t
                 )
 
+
+                unet_inputs = {
+                    'x': latent_model_input,
+                    'timesteps': torch.tensor([t] * actual_num_frames * multiplier, dtype=latent_model_input.dtype, device=device),
+                    'context': torch.cat([prompt_embeds_neg] * actual_num_frames + [prompt_embeds_pos] * actual_num_frames),
+                    'num_frames': actual_num_frames,
+                    'camera': torch.cat([camera] * multiplier),
+                }
+
+                if image is not None:
+                    unet_inputs['ip'] = torch.cat([image_embeds_neg] * actual_num_frames + [image_embeds_pos] * actual_num_frames)
+                    unet_inputs['ip_img'] = torch.cat([image_latents_neg] + [image_latents_pos]) # no repeat
+                
                 # predict the noise residual
-                noise_pred = self.unet.forward(
-                    x=latent_model_input,
-                    timesteps=torch.tensor(
-                        [t] * 4 * multiplier,
-                        dtype=latent_model_input.dtype,
-                        device=device,
-                    ),
-                    context=torch.cat(
-                        [prompt_embeds_neg] * 4 + [prompt_embeds_pos] * 4
-                    ),
-                    num_frames=4,
-                    camera=torch.cat([camera] * multiplier),
-                )
+                noise_pred = self.unet.forward(**unet_inputs)
 
                 # perform guidance
                 if do_classifier_free_guidance:
@@ -537,7 +536,6 @@ class MVDreamPipeline(DiffusionPipeline):
                     )
 
                 # compute the previous noisy sample x_t -> x_t-1
-                # latents = self.scheduler.step(noise_pred.to(dtype=torch.float32), t, latents.to(dtype=torch.float32)).prev_sample.to(prompt_embeds.dtype)
                 latents: torch.Tensor = self.scheduler.step(
                     noise_pred, t, latents, **extra_step_kwargs, return_dict=False
                 )[0]
@@ -556,7 +554,7 @@ class MVDreamPipeline(DiffusionPipeline):
         elif output_type == "pil":
             image = self.decode_latents(latents)
             image = self.numpy_to_pil(image)
-        else:
+        else: # numpy
             image = self.decode_latents(latents)
 
         # Offload last model to CPU

mvdream/util.py

@@ -1,8 +1,32 @@
 import math
 import torch
 import torch.nn as nn
+import numpy as np
 from einops import repeat
 
+from kiui.cam import orbit_camera
+
+def get_camera(
+    num_frames, elevation=15, azimuth_start=0, azimuth_span=360, blender_coord=True, extra_view=False,
+):
+    angle_gap = azimuth_span / num_frames
+    cameras = []
+    for azimuth in np.arange(azimuth_start, azimuth_span + azimuth_start, angle_gap):
+        
+        pose = orbit_camera(-elevation, azimuth, radius=1) # kiui's elevation is negated, [4, 4]
+
+        # opengl to blender
+        if blender_coord:
+            pose[2] *= -1
+            pose[[1, 2]] = pose[[2, 1]]
+
+        cameras.append(pose.flatten())
+
+    if extra_view:
+        cameras.append(np.zeros_like(cameras[0]))
+
+    return torch.from_numpy(np.stack(cameras, axis=0)).float() # [num_frames, 16]
+
 
 def checkpoint(func, inputs, params, flag):
     """

run_imagedream.py

@@ -0,0 +1,32 @@
+import torch
+import kiui
+import numpy as np
+import argparse
+from mvdream.pipeline_mvdream import MVDreamPipeline
+
+pipe = MVDreamPipeline.from_pretrained(
+    "./weights_imagedream", # local weights
+    # "ashawkey/mvdream-sd2.1-diffusers",
+    torch_dtype=torch.float16
+)
+pipe = pipe.to("cuda")
+
+
+parser = argparse.ArgumentParser(description="ImageDream")
+parser.add_argument("image", type=str, default='data/anya_rgba.png')
+parser.add_argument("--prompt", type=str, default="")
+args = parser.parse_args()
+
+while True:
+    input_image = kiui.read_image(args.image, mode='float')
+    image = pipe(args.prompt, input_image)
+    grid = np.concatenate(
+        [
+            np.concatenate([image[0], image[2]], axis=0),
+            np.concatenate([image[1], image[3]], axis=0),
+        ],
+        axis=1,
+    )
+    # kiui.vis.plot_image(grid)
+    kiui.write_image('test_imagedream.jpg', grid)
+    break

main.py → run_mvdream.py

@@ -25,4 +25,6 @@ while True:
         ],
         axis=1,
     )
-    kiui.vis.plot_image(grid)
+    # kiui.vis.plot_image(grid)
+    kiui.write_image('test_mvdream.jpg', grid)
+    break