From a5e49143ef92baf966f8b68ba2c349656f32e73e Mon Sep 17 00:00:00 2001
From: smthemex <138738845+smthemex@users.noreply.github.com>
Date: Sat, 5 Apr 2025 17:56:14 +0800
Subject: [PATCH] Update encoder_plus.py

---
 detail_encoder/encoder_plus.py | 19 +++++++++++--------
 1 file changed, 11 insertions(+), 8 deletions(-)

diff --git a/detail_encoder/encoder_plus.py b/detail_encoder/encoder_plus.py
index b922972..7f87ca9 100644
--- a/detail_encoder/encoder_plus.py
+++ b/detail_encoder/encoder_plus.py
@@ -1,13 +1,12 @@
-from typing import List
+
 import torch
-from torchvision import transforms
-from transformers import CLIPImageProcessor
+from transformers import CLIPImageProcessor,CLIPVisionConfig
 from transformers import CLIPVisionModel as OriginalCLIPVisionModel
 from ._clip import CLIPVisionModel
 from PIL import Image
 import torch.nn.functional as F
-import torch.nn as nn
-import os
+
+from safetensors.torch import load_file
 
 def is_torch2_available():
     return hasattr(F, "scaled_dot_product_attention")
@@ -19,18 +18,21 @@ from .resampler import Resampler
 
 class detail_encoder(torch.nn.Module):
     """from SSR-encoder"""
-    def __init__(self, unet, image_encoder_path, device="cuda", dtype=torch.float32):
+    def __init__(self, unet, image_encoder_path,repo, device="cuda", dtype=torch.float32):
         super().__init__()
         self.device = device
         self.dtype = dtype
 
         # load image encoder
-        clip_encoder = OriginalCLIPVisionModel.from_pretrained(image_encoder_path,attn_implementation="eager")
+        text_encoder_config = CLIPVisionConfig.from_pretrained(repo,attn_implementation="eager",local_files_only=True)
+        clip_encoder=OriginalCLIPVisionModel(text_encoder_config)
+        text_encoder_sd = load_file(image_encoder_path)
+        clip_encoder.load_state_dict(text_encoder_sd, strict=False)
         self.image_encoder = CLIPVisionModel(clip_encoder.config)
         state_dict = clip_encoder.state_dict()
         self.image_encoder.load_state_dict(state_dict, strict=False)
         self.image_encoder.to(self.device, self.dtype)
-        del clip_encoder
+        del clip_encoder,text_encoder_sd
         self.clip_image_processor = CLIPImageProcessor()
 
         # load SSR layers
@@ -111,3 +113,4 @@ class detail_encoder(torch.nn.Module):
         ).images[0]
 
         return image
+