File size: 1,890 Bytes
e8aa256 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 |
# Copyright 2023 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import Optional, Union
import torch
from torch import nn
from ...configuration_utils import ConfigMixin, register_to_config
from ...models.modeling_utils import ModelMixin
class StableUnCLIPImageNormalizer(ModelMixin, ConfigMixin):
"""
This class is used to hold the mean and standard deviation of the CLIP embedder used in stable unCLIP.
It is used to normalize the image embeddings before the noise is applied and un-normalize the noised image
embeddings.
"""
@register_to_config
def __init__(
self,
embedding_dim: int = 768,
):
super().__init__()
self.mean = nn.Parameter(torch.zeros(1, embedding_dim))
self.std = nn.Parameter(torch.ones(1, embedding_dim))
def to(
self,
torch_device: Optional[Union[str, torch.device]] = None,
torch_dtype: Optional[torch.dtype] = None,
):
self.mean = nn.Parameter(self.mean.to(torch_device).to(torch_dtype))
self.std = nn.Parameter(self.std.to(torch_device).to(torch_dtype))
return self
def scale(self, embeds):
embeds = (embeds - self.mean) * 1.0 / self.std
return embeds
def unscale(self, embeds):
embeds = (embeds * self.std) + self.mean
return embeds
|