Function compose

src/diffusers_composenW.py:44–164

(paths, category, outpath, pretrained_model_path, regularization_prompt, prompts, save_path, device='cuda')

def compose(paths, category, outpath, pretrained_model_path, regularization_prompt, prompts, save_path, device='cuda'):
    model_id = pretrained_model_path
    pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16).to("cuda")

    layers_modified = []
    for name, param in pipe.unet.named_parameters():
        if 'attn2.to_k' in name or 'attn2.to_v' in name:
            layers_modified.append(name)

    tokenizer = pipe.tokenizer

    def get_text_embedding(prompts):
        with torch.no_grad():
            uc = []
            for text in prompts:
                tokens = tokenizer(text,
                                   truncation=True,
                                   max_length=tokenizer.model_max_length,
                                   return_length=True,
                                   return_overflowing_tokens=False,
                                   padding="do_not_pad",
                                   ).input_ids
                if 'photo of a' in text[:15]:
                    print(text)
                    uc.append(pipe.text_encoder(torch.cuda.LongTensor(tokens).reshape(1,-1))[0][:, 4:].reshape(-1, 768))
                else:
                    uc.append(pipe.text_encoder(torch.cuda.LongTensor(tokens).reshape(1,-1))[0][:, 1:].reshape(-1, 768))
        return torch.cat(uc, 0).float()

    embeds = {}
    count = 1
    model2_sts = []
    modifier_tokens = []
    modifier_token_ids = []
    categories = []
    for path1, cat1 in zip(paths.split('+'), category.split('+')):
        model2_st = torch.load(path1)
        if 'modifier_token' in model2_st:
            # composition of models with individual concept only
            key = list(model2_st['modifier_token'].keys())[0]
            _ = tokenizer.add_tokens(f'<new{count}>')
            modifier_token_ids.append(tokenizer.convert_tokens_to_ids(f'<new{count}>'))
            modifier_tokens.append(True)
            embeds[f'<new{count}>'] = model2_st['modifier_token'][key]
        else:
            modifier_tokens.append(False)

        model2_sts.append(model2_st['unet'])
        categories.append(cat1)
        count += 1

    pipe.text_encoder.resize_token_embeddings(len(tokenizer))
    token_embeds = pipe.text_encoder.get_input_embeddings().weight.data
    for (x, y) in zip(modifier_token_ids, list(embeds.keys())):
        token_embeds[x] = embeds[y]
        print(x, y, "added embeddings")

    f = open(regularization_prompt, 'r')