Update app.py

app.py

@@ -355,10 +355,10 @@ def train_model(epochs):
             print(f"Original Prompt Tokens Type: {original_prompt_tokens.dtype}, Shape: {original_prompt_tokens.shape}")
             print(f"Enhanced Prompt Tokens Type: {enhanced_prompt_tokens.dtype}, Shape: {enhanced_prompt_tokens.shape}")
             # Compute prompt guidance loss (L2 norm between original and enhanced prompt embeddings)
-            prompt_loss = torch.norm(original_prompt_tokens - enhanced_prompt_tokens, p=2)
+            #prompt_loss = torch.norm(original_prompt_tokens - enhanced_prompt_tokens, p=2)
             
             # Combine losses
-            total_loss = img_loss + 0.1 * prompt_loss  # Weight the prompt guidance loss with 0.1 to balance
+            total_loss = img_loss #+ 0.1 * prompt_loss  # Weight the prompt guidance loss with 0.1 to balance
             total_loss.backward()
         
             # Optimizer step