Update safety_checker/StableDiffusionSafetyChecker.py

safety_checker/StableDiffusionSafetyChecker.py

@@ -7,8 +7,10 @@ class StableDiffusionSafetyChecker(PreTrainedModel):
 
     def __init__(self, config: CLIPConfig):
         super().__init__(config)
-        self.vision_model = CLIPVisionModel(config)
-        self.visual_projection = nn.Linear(config.hidden_size, config.projection_dim, bias=False)
+        # Truy cập hidden_size từ vision_config thay vì từ config gốc
+        self.vision_model = CLIPVisionModel(config.vision_config)
+        self.visual_projection = nn.Linear(config.vision_config.hidden_size, config.projection_dim, bias=False)
+        
         self.register_buffer("concept_embeds", torch.ones(1, 17, config.projection_dim))
         self.register_buffer("special_care_embeds", torch.ones(1, 3, config.projection_dim))
         self.register_buffer("concept_embeds_weights", torch.ones(1, 17))
@@ -16,10 +18,11 @@ class StableDiffusionSafetyChecker(PreTrainedModel):
 
     @torch.no_grad()
     def forward(self, clip_input, images):
-        pooled_output = self.vision_model(clip_input)[1]
+        pooled_output = self.vision_model(clip_input)[1]  # Get pooled_output
         image_embeds = self.visual_projection(pooled_output)
         image_embeds = image_embeds / image_embeds.norm(p=2, dim=-1, keepdim=True)
 
+        # Logic lọc nội dung
         special_cos_dist = torch.mm(image_embeds, self.special_care_embeds[0].t())
         cos_dist = torch.mm(image_embeds, self.concept_embeds[0].t())
 
@@ -28,6 +31,6 @@ class StableDiffusionSafetyChecker(PreTrainedModel):
             concept_idx = (cos_dist[i] > self.concept_embeds_weights[0]).any().item()
             has_nsfw_concepts.append(concept_idx)
             if concept_idx:
-                images[i] = torch.zeros_like(images[i])
+                images[i] = torch.zeros_like(images[i]) # Trả về ảnh đen nếu dính NSFW
         
         return images, has_nsfw_concepts