| from typing import List, Union |
|
|
| import torch |
| from PIL import Image |
| from transformers import ( |
| CLIPProcessor, |
| CLIPTextModelWithProjection, |
| CLIPTokenizer, |
| CLIPVisionModelWithProjection, |
| ) |
|
|
| from diffusers import StableDiffusionPipeline |
| from .lora import patch_pipe, tune_lora_scale, _text_lora_path, _ti_lora_path |
| import os |
| import glob |
| import math |
|
|
| EXAMPLE_PROMPTS = [ |
| "<obj> swimming in a pool", |
| "<obj> at a beach with a view of seashore", |
| "<obj> in times square", |
| "<obj> wearing sunglasses", |
| "<obj> in a construction outfit", |
| "<obj> playing with a ball", |
| "<obj> wearing headphones", |
| "<obj> oil painting ghibli inspired", |
| "<obj> working on the laptop", |
| "<obj> with mountains and sunset in background", |
| "Painting of <obj> at a beach by artist claude monet", |
| "<obj> digital painting 3d render geometric style", |
| "A screaming <obj>", |
| "A depressed <obj>", |
| "A sleeping <obj>", |
| "A sad <obj>", |
| "A joyous <obj>", |
| "A frowning <obj>", |
| "A sculpture of <obj>", |
| "<obj> near a pool", |
| "<obj> at a beach with a view of seashore", |
| "<obj> in a garden", |
| "<obj> in grand canyon", |
| "<obj> floating in ocean", |
| "<obj> and an armchair", |
| "A maple tree on the side of <obj>", |
| "<obj> and an orange sofa", |
| "<obj> with chocolate cake on it", |
| "<obj> with a vase of rose flowers on it", |
| "A digital illustration of <obj>", |
| "Georgia O'Keeffe style <obj> painting", |
| "A watercolor painting of <obj> on a beach", |
| ] |
|
|
|
|
| def image_grid(_imgs, rows=None, cols=None): |
|
|
| if rows is None and cols is None: |
| rows = cols = math.ceil(len(_imgs) ** 0.5) |
|
|
| if rows is None: |
| rows = math.ceil(len(_imgs) / cols) |
| if cols is None: |
| cols = math.ceil(len(_imgs) / rows) |
|
|
| w, h = _imgs[0].size |
| grid = Image.new("RGB", size=(cols * w, rows * h)) |
| grid_w, grid_h = grid.size |
|
|
| for i, img in enumerate(_imgs): |
| grid.paste(img, box=(i % cols * w, i // cols * h)) |
| return grid |
|
|
|
|
| def text_img_alignment(img_embeds, text_embeds, target_img_embeds): |
| |
| |
|
|
| |
| assert img_embeds.shape[0] == text_embeds.shape[0] |
| text_img_sim = (img_embeds * text_embeds).sum(dim=-1) / ( |
| img_embeds.norm(dim=-1) * text_embeds.norm(dim=-1) |
| ) |
|
|
| |
| img_embed_normalized = img_embeds / img_embeds.norm(dim=-1, keepdim=True) |
|
|
| avg_target_img_embed = ( |
| (target_img_embeds / target_img_embeds.norm(dim=-1, keepdim=True)) |
| .mean(dim=0) |
| .unsqueeze(0) |
| .repeat(img_embeds.shape[0], 1) |
| ) |
|
|
| img_img_sim = (img_embed_normalized * avg_target_img_embed).sum(dim=-1) |
|
|
| return { |
| "text_alignment_avg": text_img_sim.mean().item(), |
| "image_alignment_avg": img_img_sim.mean().item(), |
| "text_alignment_all": text_img_sim.tolist(), |
| "image_alignment_all": img_img_sim.tolist(), |
| } |
|
|
|
|
| def prepare_clip_model_sets(eval_clip_id: str = "openai/clip-vit-large-patch14"): |
| text_model = CLIPTextModelWithProjection.from_pretrained(eval_clip_id) |
| tokenizer = CLIPTokenizer.from_pretrained(eval_clip_id) |
| vis_model = CLIPVisionModelWithProjection.from_pretrained(eval_clip_id) |
| processor = CLIPProcessor.from_pretrained(eval_clip_id) |
|
|
| return text_model, tokenizer, vis_model, processor |
|
|
|
|
| def evaluate_pipe( |
| pipe, |
| target_images: List[Image.Image], |
| class_token: str = "", |
| learnt_token: str = "", |
| guidance_scale: float = 5.0, |
| seed=0, |
| clip_model_sets=None, |
| eval_clip_id: str = "openai/clip-vit-large-patch14", |
| n_test: int = 10, |
| n_step: int = 50, |
| ): |
|
|
| if clip_model_sets is not None: |
| text_model, tokenizer, vis_model, processor = clip_model_sets |
| else: |
| text_model, tokenizer, vis_model, processor = prepare_clip_model_sets( |
| eval_clip_id |
| ) |
|
|
| images = [] |
| img_embeds = [] |
| text_embeds = [] |
| for prompt in EXAMPLE_PROMPTS[:n_test]: |
| prompt = prompt.replace("<obj>", learnt_token) |
| torch.manual_seed(seed) |
| with torch.autocast("cuda"): |
| img = pipe( |
| prompt, num_inference_steps=n_step, guidance_scale=guidance_scale |
| ).images[0] |
| images.append(img) |
|
|
| |
| inputs = processor(images=img, return_tensors="pt") |
| img_embed = vis_model(**inputs).image_embeds |
| img_embeds.append(img_embed) |
|
|
| prompt = prompt.replace(learnt_token, class_token) |
| |
| inputs = tokenizer([prompt], padding=True, return_tensors="pt") |
| outputs = text_model(**inputs) |
| text_embed = outputs.text_embeds |
| text_embeds.append(text_embed) |
|
|
| |
| inputs = processor(images=target_images, return_tensors="pt") |
| target_img_embeds = vis_model(**inputs).image_embeds |
|
|
| img_embeds = torch.cat(img_embeds, dim=0) |
| text_embeds = torch.cat(text_embeds, dim=0) |
|
|
| return text_img_alignment(img_embeds, text_embeds, target_img_embeds) |
|
|
|
|
| def visualize_progress( |
| path_alls: Union[str, List[str]], |
| prompt: str, |
| model_id: str = "runwayml/stable-diffusion-v1-5", |
| device="cuda:0", |
| patch_unet=True, |
| patch_text=True, |
| patch_ti=True, |
| unet_scale=1.0, |
| text_sclae=1.0, |
| num_inference_steps=50, |
| guidance_scale=5.0, |
| offset: int = 0, |
| limit: int = 10, |
| seed: int = 0, |
| ): |
|
|
| imgs = [] |
| if isinstance(path_alls, str): |
| alls = list(set(glob.glob(path_alls))) |
|
|
| alls.sort(key=os.path.getmtime) |
| else: |
| alls = path_alls |
|
|
| pipe = StableDiffusionPipeline.from_pretrained( |
| model_id, torch_dtype=torch.float16 |
| ).to(device) |
|
|
| print(f"Found {len(alls)} checkpoints") |
| for path in alls[offset:limit]: |
| print(path) |
|
|
| patch_pipe( |
| pipe, path, patch_unet=patch_unet, patch_text=patch_text, patch_ti=patch_ti |
| ) |
|
|
| tune_lora_scale(pipe.unet, unet_scale) |
| tune_lora_scale(pipe.text_encoder, text_sclae) |
|
|
| torch.manual_seed(seed) |
| image = pipe( |
| prompt, |
| num_inference_steps=num_inference_steps, |
| guidance_scale=guidance_scale, |
| ).images[0] |
| imgs.append(image) |
|
|
| return imgs |
|
|
