Update app.py

app.py

@@ -172,9 +172,14 @@ def loss_style(prompt, style_embed, style_seed):
     # Assuming token_embeddings has shape [batch_size, seq_length, embedding_dim]
     replacement_token_embedding = replacement_token_embedding[:768]  # Adjust the size
     replacement_token_embedding = replacement_token_embedding.unsqueeze(0)  # Make it [1, 768] if necessary
-
+    indices = torch.where(input_ids[0] == 6829)[0]  # Extract indices where the condition is True
+    print(f"indices: {indices}")  # Debug print
+    for index in indices:
+        print(f"index: {index}")  # Debug print
+        token_embeddings[0, index] = replacement_token_embedding.to(torch_device)  # Update each index
+    
     # Insert this into the token embeddings
-    token_embeddings[0, torch.where(input_ids[0]==6829)] = replacement_token_embedding.to(torch_device)
+   # token_embeddings[0, torch.where(input_ids[0]==6829)] = replacement_token_embedding.to(torch_device)
 
     # Combine with pos embs
     input_embeddings = token_embeddings + position_embeddings