Update app.py

app.py

@@ -342,14 +342,14 @@ def train_model(epochs):
         for i, (original, target, original_prompt_tokens, enhanced_prompt_tokens) in enumerate(dataloader):
             # Move images and prompt embeddings to the appropriate device (CPU or GPU)
             original, target = original.to(device), target.to(device)
-            original_prompt_tokens = original_prompt_tokens.input_ids.to(device)
-            enhanced_prompt_tokens = enhanced_prompt_tokens.input_ids.to(device)
+            original_prompt_tokens = original_prompt_tokens.input_ids.to(device).float()  # Convert to float
+            enhanced_prompt_tokens = enhanced_prompt_tokens.input_ids.to(device).float()  # Convert to float
             
             optimizer.zero_grad()
-
+        
             # Forward pass through the model
             output = model(target)
-
+        
             # Compute image reconstruction loss
             img_loss = criterion(output, original)
             
@@ -359,7 +359,7 @@ def train_model(epochs):
             # Combine losses
             total_loss = img_loss + 0.1 * prompt_loss  # Weight the prompt guidance loss with 0.1 to balance
             total_loss.backward()
-
+        
             # Optimizer step
             optimizer.step()