Update app.py

app.py

@@ -23,7 +23,6 @@ from taming.models.vqgan import VQModel
 from omegaconf import OmegaConf
 from taming.models.vqgan import GumbelVQ
 import gradio as gr
-from modules.finetunedvqgan import Generator
 from modules.modelz import DeepfakeToSourceTransformer
 from modules.frameworkeval import DF
 from modules.segmentface import FaceSegmenter
@@ -31,6 +30,25 @@ from modules.denormalize import denormalize_bin, denormalize_tr, denormalize_ar
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
+config = OmegaConf.load(self.config_path)
+# Extract parameters specific to GumbelVQ
+vq_params = config.model.params
+# Initialize the GumbelVQ models
+model_vaq_f = GumbelVQ(
+            ddconfig=vq_params.ddconfig,
+            lossconfig=vq_params.lossconfig,
+            n_embed=vq_params.n_embed,
+            embed_dim=vq_params.embed_dim,
+            kl_weight=vq_params.kl_weight,
+            temperature_scheduler_config=vq_params.temperature_scheduler_config).to(device)
+model_vaq_g = GumbelVQ(
+            ddconfig=vq_params.ddconfig,
+            lossconfig=vq_params.lossconfig,
+            n_embed=vq_params.n_embed,
+            embed_dim=vq_params.embed_dim,
+            kl_weight=vq_params.kl_weight,
+            temperature_scheduler_config=vq_params.temperature_scheduler_config).to(device)
+
 ##________________________Transformation______________________________
 
 transform = T.Compose([
@@ -54,10 +72,17 @@ def gen_sources(deepfake_img):
     #------------Initialize:Decoder-F------------------------
     config_path = "./models/config.yaml"
     checkpoint_path_f = "./models/model_vaq1_ff.pth"
-    model_vaq_f = Generator(config_path, checkpoint_path_f, device)
+    # Load model checkpoints
+    checkpoint_f = torch.load(self.checkpoint_path_f, map_location=self.device)
+    # Load the state dictionary into the models
+    model_vaq_f = model_vaq_f.load_state_dict(checkpoint_f, strict=True)
+    model_vaq_f.eval()
     #------------Initialize:Decoder-G------------------------
     checkpoint_path_g = "./models/model_vaq2_gg.pth"
-    model_vaq_g = Generator(config_path, checkpoint_path_g, device)
+    checkpoint_g = torch.load(self.checkpoint_path_g, map_location=self.device)
+    # Load the state dictionary into the models
+    model_vaq_g = model_vaq_g.load_state_dict(checkpoint_g, strict=True)
+    model_vaq_g.eval()
     ##------------------------Initialize Model-F-------------------------------------
     model_z1 = DeepfakeToSourceTransformer().to(device)
     model_z1.load_state_dict(torch.load("./models/model_z1_ff.pth",map_location=device),strict=True)
@@ -73,7 +98,6 @@ def gen_sources(deepfake_img):
     with torch.no_grad():
         # Load and preprocess input image
         #img = Image.open(deepfake_img).convert('RGB')
-        #segimg = Image.open(deepfake_seg).convert('RGB')
         df_img = transform(deepfake_img.convert('RGB')).unsqueeze(0).to(device)  # Shape: (1, 3, 256, 256)
         seg_img = transform(deepfake_seg).unsqueeze(0).to(device)