Update app.py

app.py

@@ -101,8 +101,9 @@ def gen_sources(deepfake_img):
         rec_z_img2 = model_z2(z_seg)
         rec_img1 = model_vaq_f.decode(rec_z_img1).squeeze(0)
         rec_img2 = model_vaq_g.decode(rec_z_img2).squeeze(0)
-        rec_img1_pil = T.ToPILImage()(rec_img1)
-        rec_img2_pil = T.ToPILImage()(rec_img2)
+        rec_img1_pil = T.ToPILImage()(denormalize_bin(rec_img1))
+        rec_img2_pil = T.ToPILImage()(denormalize_bin(rec_img2))
+
 
         # Save PIL images to temporary files
         with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as temp1, \
@@ -129,7 +130,7 @@ def gen_sources(deepfake_img):
         dfimage_pil = Image.open(result)
         with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as temp3:
             dfimage_pil.save(temp3, format="PNG")
-            rec_df = transform(Image.open(temp3.name)).unsqueeze(0).to(device)
+            rec_df = denormalize_bin(transform(Image.open(temp3.name))).unsqueeze(0).to(device)
             os.remove(temp3.name)
 
         rec_loss, _ = criterion(df_img, rec_df)