|
import torch |
|
import torch.nn as nn |
|
from .config import CSDConfig |
|
from transformers import PreTrainedModel, CLIPVisionModel |
|
|
|
class CSDModel(PreTrainedModel): |
|
config_class = CSDConfig |
|
def __init__(self, config: CSDConfig): |
|
super().__init__(config) |
|
self.backbone = CLIPVisionModel(config) |
|
self.out_style = nn.Linear(config.hidden_size, config.style_projection_dim, bias=False) |
|
self.out_content = nn.Linear(config.hidden_size, config.content_projection_dim, bias=False) |
|
|
|
@torch.inference_mode() |
|
def forward(self, pixel_values): |
|
features = self.backbone(pixel_values, return_dict=False)[1] |
|
style_embeds = self.out_style(features) |
|
content_embeds = self.out_content(features) |
|
return features, style_embeds, content_embeds |
