Update app.py

app.py

@@ -86,8 +86,8 @@ class Pix2PixDataset(torch.utils.data.Dataset):
         enhanced_prompt = self.data.iloc[idx]['enhanced_prompt']
 
         # Encode images
-        original_image_latents = vae.encode(original_images).latent_dist.sample()
-        target_image_latents = vae.encode(target_images).latent_dist.sample()
+        original_image_latents = vae.encode(original_image).latent_dist.sample()
+        target_image_latents = vae.encode(target_image).latent_dist.sample()
         
         # Encode prompts
         prompt_latents = encode_prompt(enhanced_prompt,clip_model,clip_tokenizer)