draft imagedream, before merge

README.md

@@ -8,11 +8,22 @@ modified from https://github.com/KokeCacao/mvdream-hf.
 pip install -U omegaconf diffusers safetensors huggingface_hub transformers accelerate
 
 # download original ckpt
+cd models
 wget https://huggingface.co/MVDream/MVDream/resolve/main/sd-v2.1-base-4view.pt
 wget https://raw.githubusercontent.com/bytedance/MVDream/main/mvdream/configs/sd-v2-base.yaml
+cd ..
 
 # convert
-python convert_mvdream_to_diffusers.py --checkpoint_path ./sd-v2.1-base-4view.pt --dump_path ./weights --original_config_file ./sd-v2-base.yaml --half --to_safetensors --test
+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
+```
+
+```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
 ```
 
 ### usage

convert_imagedream_to_diffusers.py

@@ -0,0 +1,561 @@
+# 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!")

imagedream/adaptor.py

@@ -0,0 +1,141 @@
+import math
+import torch
+import torch.nn as nn
+
+# FFN
+def FeedForward(dim, mult=4):
+    inner_dim = int(dim * mult)
+    return nn.Sequential(
+        nn.LayerNorm(dim),
+        nn.Linear(dim, inner_dim, bias=False),
+        nn.GELU(),
+        nn.Linear(inner_dim, dim, bias=False),
+    )
+
+
+def reshape_tensor(x, heads):
+    bs, length, width = x.shape
+    # (bs, length, width) --> (bs, length, n_heads, dim_per_head)
+    x = x.view(bs, length, heads, -1)
+    # (bs, length, n_heads, dim_per_head) --> (bs, n_heads, length, dim_per_head)
+    x = x.transpose(1, 2)
+    # (bs, n_heads, length, dim_per_head) --> (bs*n_heads, length, dim_per_head)
+    x = x.reshape(bs, heads, length, -1)
+    return x
+
+
+class PerceiverAttention(nn.Module):
+    def __init__(self, *, dim, dim_head=64, heads=8):
+        super().__init__()
+        self.scale = dim_head ** -0.5
+        self.dim_head = dim_head
+        self.heads = heads
+        inner_dim = dim_head * heads
+
+        self.norm1 = nn.LayerNorm(dim)
+        self.norm2 = nn.LayerNorm(dim)
+
+        self.to_q = nn.Linear(dim, inner_dim, bias=False)
+        self.to_kv = nn.Linear(dim, inner_dim * 2, bias=False)
+        self.to_out = nn.Linear(inner_dim, dim, bias=False)
+
+    def forward(self, x, latents):
+        """
+        Args:
+            x (torch.Tensor): image features
+                shape (b, n1, D)
+            latent (torch.Tensor): latent features
+                shape (b, n2, D)
+        """
+        x = self.norm1(x)
+        latents = self.norm2(latents)
+
+        b, l, _ = latents.shape
+
+        q = self.to_q(latents)
+        kv_input = torch.cat((x, latents), dim=-2)
+        k, v = self.to_kv(kv_input).chunk(2, dim=-1)
+
+        q = reshape_tensor(q, self.heads)
+        k = reshape_tensor(k, self.heads)
+        v = reshape_tensor(v, self.heads)
+
+        # attention
+        scale = 1 / math.sqrt(math.sqrt(self.dim_head))
+        weight = (q * scale) @ (k * scale).transpose(
+            -2, -1
+        )  # More stable with f16 than dividing afterwards
+        weight = torch.softmax(weight.float(), dim=-1).type(weight.dtype)
+        out = weight @ v
+
+        out = out.permute(0, 2, 1, 3).reshape(b, l, -1)
+
+        return self.to_out(out)
+
+
+class ImageProjModel(torch.nn.Module):
+    """Projection Model"""
+
+    def __init__(
+        self,
+        cross_attention_dim=1024,
+        clip_embeddings_dim=1024,
+        clip_extra_context_tokens=4,
+    ):
+        super().__init__()
+        self.cross_attention_dim = cross_attention_dim
+        self.clip_extra_context_tokens = clip_extra_context_tokens
+
+        # from 1024 -> 4 * 1024
+        self.proj = torch.nn.Linear(
+            clip_embeddings_dim, self.clip_extra_context_tokens * cross_attention_dim
+        )
+        self.norm = torch.nn.LayerNorm(cross_attention_dim)
+
+    def forward(self, image_embeds):
+        embeds = image_embeds
+        clip_extra_context_tokens = self.proj(embeds).reshape(
+            -1, self.clip_extra_context_tokens, self.cross_attention_dim
+        )
+        clip_extra_context_tokens = self.norm(clip_extra_context_tokens)
+        return clip_extra_context_tokens
+
+
+class Resampler(nn.Module):
+    def __init__(
+        self,
+        dim=1024,
+        depth=8,
+        dim_head=64,
+        heads=16,
+        num_queries=8,
+        embedding_dim=768,
+        output_dim=1024,
+        ff_mult=4,
+    ):
+        super().__init__()
+        self.latents = nn.Parameter(torch.randn(1, num_queries, dim) / dim ** 0.5)
+        self.proj_in = nn.Linear(embedding_dim, dim)
+        self.proj_out = nn.Linear(dim, output_dim)
+        self.norm_out = nn.LayerNorm(output_dim)
+
+        self.layers = nn.ModuleList([])
+        for _ in range(depth):
+            self.layers.append(
+                nn.ModuleList(
+                    [
+                        PerceiverAttention(dim=dim, dim_head=dim_head, heads=heads),
+                        FeedForward(dim=dim, mult=ff_mult),
+                    ]
+                )
+            )
+
+    def forward(self, x):
+        latents = self.latents.repeat(x.size(0), 1, 1)
+        x = self.proj_in(x)
+        for attn, ff in self.layers:
+            latents = attn(x, latents) + latents
+            latents = ff(latents) + latents
+
+        latents = self.proj_out(latents)
+        return self.norm_out(latents)

imagedream/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,35 @@ 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,
+            with_ip=False,
+            ip_dim=16,
+            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.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)
@@ -128,9 +85,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.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)
 
@@ -149,8 +115,21 @@ class MemoryEfficientCrossAttention(nn.Module):
             q, k, v, attn_bias=None, op=self.attention_op
         )
 
-        if mask is not None:
-            raise NotImplementedError
+        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)
@@ -160,148 +139,47 @@ 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,
+        context_dim,
         n_heads,
         d_head,
         dropout=0.0,
-        context_dim=None,
         gated_ff=True,
         checkpoint=True,
-        disable_self_attn=False,
+        with_ip=False,
+        ip_dim=16,
+        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
+            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):
-        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 +187,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 +195,32 @@ 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,
+        with_ip=False,
+        ip_dim=16,
+        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 +228,19 @@ 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,
+                    with_ip=with_ip,
+                    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 +249,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

imagedream/models.py

@@ -13,8 +13,8 @@ from .util import (
     zero_module,
     timestep_embedding,
 )
-from .attention import SpatialTransformer, SpatialTransformer3D
-
+from .attention import SpatialTransformer3D
+from .adaptor import Resampler, ImageProjModel
 
 class CondSequential(nn.Sequential):
     """
@@ -28,8 +28,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 +272,9 @@ class MultiViewUNetModel(ModelMixin, ConfigMixin):
         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__()
@@ -305,9 +306,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 +333,22 @@ class MultiViewUNetModel(ModelMixin, ConfigMixin):
         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),
@@ -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,12 @@ 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,
+                                with_ip=self.with_ip,
+                                ip_dim=self.ip_dim,
+                                ip_weight=self.ip_weight,
                             )
                         )
                 self.input_blocks.append(CondSequential(*layers))
@@ -463,10 +475,12 @@ 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,
+                with_ip=self.with_ip,
+                ip_dim=self.ip_dim,
+                ip_weight=self.ip_weight,
             ), 
             ResBlock(
                 ch,
@@ -501,11 +515,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 +522,12 @@ 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,
+                                with_ip=self.with_ip,
+                                ip_dim=self.ip_dim,
+                                ip_weight=self.ip_weight,
                             )
                         )
                 if level and i == self.num_res_blocks[level]:
@@ -559,6 +570,9 @@ class MultiViewUNetModel(ModelMixin, ConfigMixin):
         y: Optional[Tensor] = None,
         camera=None,
         num_frames=1,
+        # should be provided if with_ip
+        ip = None,
+        ip_img = None,
         **kwargs,
     ):
         """
@@ -592,6 +606,11 @@ class MultiViewUNetModel(ModelMixin, ConfigMixin):
         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:

imagedream/pipeline_imagedream.py

@@ -2,7 +2,7 @@ import torch
 import inspect
 import numpy as np
 from typing import Callable, List, Optional, Union
-from transformers import CLIPTextModel, CLIPTokenizer, CLIPVisionModel, CLIPImageProcessor
+from transformers import CLIPTextModel, CLIPTokenizer, CLIPVisionModel, CLIPFeatureExtractor
 from diffusers import AutoencoderKL, DiffusionPipeline
 from diffusers.utils import (
     deprecate,
@@ -16,6 +16,8 @@ from diffusers.utils.torch_utils import randn_tensor
 
 from .models import MultiViewUNetModel
 
+import kiui
+
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
 
 
@@ -62,7 +64,7 @@ def convert_opengl_to_blender(camera_matrix):
 
 
 def get_camera(
-    num_frames, elevation=15, azimuth_start=0, azimuth_span=360, blender_coord=True
+    num_frames, elevation=15, azimuth_start=0, azimuth_span=360, blender_coord=True, extra_view=False,
 ):
     angle_gap = azimuth_span / num_frames
     cameras = []
@@ -71,6 +73,9 @@ def get_camera(
         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()
 
 
@@ -82,8 +87,8 @@ class ImageDreamPipeline(DiffusionPipeline):
         tokenizer: CLIPTokenizer,
         text_encoder: CLIPTextModel,
         scheduler: DDIMScheduler,
-        feature_extractor: CLIPImageProcessor,
-        image_encoder: CLIPVisionModel,
+        feature_extractor: CLIPFeatureExtractor = None,
+        image_encoder: CLIPVisionModel = None,
         requires_safety_checker: bool = False,
     ):
         super().__init__()
@@ -449,10 +454,36 @@ class ImageDreamPipeline(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).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 (TODO: pil?)
+        image, # input image, np.ndarray float32!
         prompt: str = "a car",
         height: int = 256,
         width: int = 256,
@@ -465,7 +496,7 @@ class ImageDreamPipeline(DiffusionPipeline):
         output_type: Optional[str] = "image",
         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)
@@ -482,7 +513,18 @@ class ImageDreamPipeline(DiffusionPipeline):
         self.scheduler.set_timesteps(num_inference_steps, device=device)
         timesteps = self.scheduler.timesteps
 
-        _prompt_embeds: torch.Tensor = self._encode_prompt(
+        # 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,
@@ -493,8 +535,8 @@ class ImageDreamPipeline(DiffusionPipeline):
 
         # Prepare latent variables
         latents: torch.Tensor = self.prepare_latents(
-            batch_size * num_images_per_prompt,
-            4,
+            (num_frames + 1) * num_images_per_prompt,
+            4, # channel
             height,
             width,
             prompt_embeds_pos.dtype,
@@ -503,9 +545,10 @@ class ImageDreamPipeline(DiffusionPipeline):
             None,
         )
 
-        camera = get_camera(batch_size).to(dtype=latents.dtype, device=device)
+        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. 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
@@ -523,15 +566,22 @@ class ImageDreamPipeline(DiffusionPipeline):
                 noise_pred = self.unet.forward(
                     x=latent_model_input,
                     timesteps=torch.tensor(
-                        [t] * 4 * multiplier,
+                        [t] * (num_frames + 1) * multiplier,
                         dtype=latent_model_input.dtype,
                         device=device,
                     ),
                     context=torch.cat(
-                        [prompt_embeds_neg] * 4 + [prompt_embeds_pos] * 4
+                        [prompt_embeds_neg] * (num_frames + 1) + [prompt_embeds_pos] * (num_frames + 1)
                     ),
-                    num_frames=4,
+                    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
@@ -542,7 +592,6 @@ class ImageDreamPipeline(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]