Class ClipImageEmbedder

diff2flow/conditioning/encoders.py:12–62

# Models in ["ViT-B/32", "ViT-B/16", "ViT-L/14", "ViT-L/14@336px"]
class ClipImageEmbedder(nn.Module):
    def __init__(
            self,
            model="ViT-L/14",
            jit=False,
            device='cuda' if torch.cuda.is_available() else 'cpu',
            antialias=True,
            ucg_rate=0.
    ):
        super().__init__()
        from clip import load as load_clip
        self.model, _ = load_clip(name=model, device=device, jit=jit)

        self.antialias = antialias

        self.register_buffer('mean', torch.Tensor([0.48145466, 0.4578275, 0.40821073]), persistent=False)
        self.register_buffer('std', torch.Tensor([0.26862954, 0.26130258, 0.27577711]), persistent=False)
        self.ucg_rate = ucg_rate

        self.init_uncond()

    def init_uncond(self, path="nulltext.npy"):
        try:
            assert os.path.exists(path), f"Uncond file {path} not found."
            print(f"Loading uncond from {path}")
            uncond = torch.from_numpy(np.load(path))
            self.register_buffer("uncond", uncond)
        except:
            self.uncond = None

    def preprocess(self, x):
        # resize to 224, normalize to [0,1] and re-normalize according to clip
        x = torch.nn.functional.interpolate(x, size=(224, 224), mode='bilinear', align_corners=False)
        assert x.min() >= -1. and x.max() <= 1.
        x = (x + 1.) / 2.
        x = (x - self.mean.reshape(1,-1,1,1)) / self.std.reshape(1,-1,1,1)
        return x

    def forward(self, x, no_dropout=False):
        # x is assumed to be in range [-1,1]
        out = self.model.encode_image(self.preprocess(x))
        out = out.to(x.dtype)
        if self.ucg_rate > 0. and not no_dropout:
            out = torch.bernoulli((1. - self.ucg_rate) * torch.ones(out.shape[0], device=out.device))[:, None] * out
        return out.unsqueeze(1)
    
    @torch.no_grad()
    def get_unconditional_conditioning(self, device="cuda"):
        if self.uncond is None:
            raise ValueError("Unconditional conditioning not initialized.")
        return self.uncond.to(device)


class DummyOpenCLIPTextEmbedder(nn.Module):