| 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 |
