Help verify StableDiffusion & CLIP weight sharing

🧨 Diffusers
1

Hi, can someone verify my approach to be able to use both StableDiffusionImg2ImgPipeline and CLIP similarity/distance without loading any weights twice?

import torch
from transformers import CLIPTextModelWithProjection
from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline
from diffusers.pipelines.stable_diffusion.safety_checker import cosine_distance
from PIL import Image

## StableDiffusionImg2Img ##
img_to_img_pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
    "./stable-diffusion-v1-5", torch_dtype=torch.float16, revision="fp16"
)
img_to_img_pipe = img_to_img_pipe.to("cuda")

## CLIP similarity ##
feature_extractor=img_to_img_pipe.feature_extractor
vision_model = img_to_img_pipe.safety_checker.vision_model
visual_projection = img_to_img_pipe.safety_checker.visual_projection
text_model = img_to_img_pipe.text_encoder
tokenizer = img_to_img_pipe.tokenizer

# Image embeddings
image = Image.open("dog.png")
image_features = feature_extractor([image], return_tensors="pt")
pixel_values = image_features.pixel_values.to(device=vision_model.device, dtype=vision_model.dtype)
pooled_output = vision_model(pixel_values)[1]
image_embeds = visual_projection(pooled_output)

# Text embeddings
prompt = "photo of a dog"
text_tokens = tokenizer(
    prompt,
    padding="max_length",
    max_length=tokenizer.model_max_length,
    truncation=True,
    return_tensors="pt",
)
text_tokens_ids = text_tokens.input_ids.to(device=text_model.device)
pooled_output = text_model(input_ids=text_tokens_ids)[1]
# Text model has no projection layer :( will extract only the text_projection in the future
text_encoder = CLIPTextModelWithProjection.from_pretrained("openai/clip-vit-large-patch14", torch_dtype=torch.float16)
text_encoder = text_encoder.to("cuda")
text_embeds = text_encoder.text_projection(pooled_output)

# Get similarity
similarity = cosine_distance(image_embeds, text_embeds)

print("Similarity:", similarity)

Info about vision model version: config.json · CompVis/stable-diffusion-safety-checker at main
Info about text model version: stable-diffusion/modules.py at main · runwayml/stable-diffusion · GitHub