Function compose

src/composenW.py:68–265

(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, config = get_model(pretrained_model_path)
    model.eval()
    model.requires_grad = False

    layers = []
    layers_modified = []

    def getlayers(model, root_name=''):
        for name, module in model.named_children():
            if module.__class__.__name__ == 'SpatialTransformer':
                layers_modified.append(root_name + '.' + name + '.transformer_blocks.0.attn2.to_k')
                layers_modified.append(root_name + '.' + name + '.transformer_blocks.0.attn2.to_v')
            else:
                if list(module.children()) == []:
                    layers.append(root_name + '.' + name)
                else:
                    getlayers(module, root_name + '.' + name)

    getlayers(model.model.diffusion_model)

    for i in range(len(layers_modified)):
        layers_modified[i] = 'model.diffusion_model' + layers_modified[i] + '.weight'

    def get_text_embedding(prompts):
        with torch.no_grad():
            uc = []
            for text in prompts:
                tokens = tokenizer(text,
                                   truncation=True,
                                   max_length=77,
                                   return_length=True,
                                   return_overflowing_tokens=False,
                                   padding="max_length",
                                   return_tensors="pt")

                tokens = tokens["input_ids"]
                end = torch.nonzero(tokens == 49407)[:, 1].min()
                if 'photo of a' in text[:15]:
                    print(text)
                    uc.append((model.get_learned_conditioning(1 * [text])[:, 4:end+1]).reshape(-1, 768))
                else:
                    uc.append((model.get_learned_conditioning(1 * [text])[:, 1:end+1]).reshape(-1, 768))

        return torch.cat(uc, 0)

    tokenizer = model.cond_stage_model.tokenizer
    embeds = []
    count = 1

    model2_sts = []
    modifier_tokens = []
    categories = []
    config.model.params.cond_stage_config.params = {}
    config.model.params.cond_stage_config.params.modifier_token = None
    for path1, cat1 in zip(paths.split('+'), category.split('+')):
        model2_st = torch.load(path1)
        if 'embed' in model2_st['state_dict']: