Update utils.py

utils.py

@@ -218,3 +218,39 @@ token_emb_layer = text_encoder.text_model.embeddings.token_embedding
 
 pos_emb_layer = text_encoder.text_model.embeddings.position_embedding
 #pos_emb_layer
+
+def func_generate(query, concept_idx, seed_start, contrast_loss=False, contrast_perc=None):
+    prompt = query + ' in the style of bulb'
+    text_input = tokenizer(prompt, padding="max_length", max_length=tokenizer.model_max_length, truncation=True,
+                           return_tensors="pt")
+    input_ids = text_input.input_ids.to(torch_device)
+
+    # Get token embeddings
+    position_ids = text_encoder.text_model.embeddings.position_ids[:, :77]
+    position_embeddings = pos_emb_layer(position_ids)
+
+    s = seed_start
+
+    token_embeddings = token_emb_layer(input_ids)
+    # The new embedding - our special birb word
+    replacement_token_embedding = concept_embeds[concept_idx].to(torch_device)
+
+    # Insert this into the token embeddings
+    token_embeddings[0, torch.where(input_ids[0] == 22373)] = replacement_token_embedding.to(torch_device)
+
+    # Combine with pos embs
+    input_embeddings = token_embeddings + position_embeddings
+
+    #  Feed through to get final output embs
+    modified_output_embeddings = get_output_embeds(input_embeddings)
+
+    # And generate an image with this:
+
+    if contrast_loss and seed_values[concept_idx] > 0:
+        s = seed_values[concept_idx]
+    else:
+        s = random.randint(s + 1, s + 30)
+        seed_values[concept_idx] = s
+    
+    g = torch.manual_seed(s)
+    return generate_with_embs(text_input, modified_output_embeddings, generator=g, contrast_loss=contrast_loss, contrast_perc=contrast_perc)