Update app.py

app.py

@@ -58,14 +58,14 @@ model_vaq_g = GumbelVQ(
 transform = T.Compose([
     T.Resize((256, 256)),   
     T.ToTensor(),         
-    T.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])])  # Normalize to [-1, 1]
+    T.Normalize(mean=[0, 0, 0], std=[1, 1, 1])])  # Normalize to [-1, 1]
 
 #_________________Define:Gradio Function________________________
 
 def gen_sources(deepfake_img):
     #----------------DeepFake Face Segmentation-----------------
     segmenter = FaceSegmenter(threshold=0.5)
-    img_np = np.array(deepfake_img)
+    img_np = np.array(deepfake_img.convert('RGB'))
     img_bgr = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
     segmented_np = segmenter.segment_face(img_bgr)
     deepfake_seg = Image.fromarray(cv2.cvtColor(segmented_np, cv2.COLOR_BGR2RGB))
@@ -92,7 +92,7 @@ def gen_sources(deepfake_img):
     criterion = DF()
 
     with torch.no_grad():
-        df_img = transform(deepfake_img).unsqueeze(0).to(device)
+        df_img = transform(deepfake_img.convert('RGB')).unsqueeze(0).to(device)
         seg_img = transform(deepfake_seg).unsqueeze(0).to(device)
 
         z_df, _, _ = model_vaq_f.encode(df_img)