Update app.py

app.py

@@ -75,9 +75,13 @@ original_generator.load_state_dict(ckpt["g_ema"], strict=False)
 mean_latent = original_generator.mean_latent(10000)
 # print(ckpt.keys())
 
-generatorjoker = deepcopy(original_generator)
+generatorjokerfalse = deepcopy(original_generator)
 
-generatorvoldemort = deepcopy(original_generator)
+generatorjokertrue = deepcopy(original_generator)
+
+generatorvoldemortfalse = deepcopy(original_generator)
+
+generatorvoldemorttrue = deepcopy(original_generator)
 
 generatorpushpa = deepcopy(original_generator)
 
@@ -103,15 +107,21 @@ transform = transforms.Compose(
 )
 
 
-modeljoker = hf_hub_download(repo_id="Abhinowww/Capstone", filename="JokerEightHundredFalse.pt")
-ckptjoker = torch.load(modeljoker, map_location=lambda storage, loc: storage)
-generatorjoker.load_state_dict(ckptjoker, strict=False)
+modeljokerfalse = hf_hub_download(repo_id="Abhinowww/Capstone", filename="JokerEightHundredFalse.pt")
+ckptjokerfalse = torch.load(modeljokerfalse, map_location=lambda storage, loc: storage)
+generatorjokerfalse.load_state_dict(ckptjokerfalse, strict=False)
 
+modeljokertrue = hf_hub_download(repo_id="Abhinowww/Capstone", filename="JokerTwoHundredFiftyTrue.pt")
+ckptjokertrue = torch.load(modeljokertrue, map_location=lambda storage, loc: storage)
+generatorjokertrue.load_state_dict(ckptjokertrue, strict=False)
 
-modelvoldemort = hf_hub_download(repo_id="Abhinowww/Capstone", filename="VoldemortFourHundredFalse.pt")
-ckptvoldemort = torch.load(modelvoldemort, map_location=lambda storage, loc: storage)
-generatorvoldemort.load_state_dict(ckptvoldemort, strict=False)
+modelvoldemortfalse = hf_hub_download(repo_id="Abhinowww/Capstone", filename="VoldemortFourHundredFalse.pt")
+ckptvoldemortfalse = torch.load(modelvoldemortfalse, map_location=lambda storage, loc: storage)
+generatorvoldemortfalse.load_state_dict(ckptvoldemortfalse, strict=False)
 
+modelvoldemorttrue = hf_hub_download(repo_id="Abhinowww/Capstone", filename="VoldemortThreeHundredTrue.pt")
+ckptvoldemorttrue = torch.load(modelvoldemorttrue, map_location=lambda storage, loc: storage)
+generatorvoldemorttrue.load_state_dict(ckptvoldemorttrue, strict=False)
 
 modelpushpa = hf_hub_download(repo_id="Abhinowww/Capstone", filename="PushpaFourHundredFalse.pt")
 ckptpushpa = torch.load(modelpushpa, map_location=lambda storage, loc: storage)
@@ -121,11 +131,11 @@ modelgiga = hf_hub_download(repo_id="Abhinowww/Capstone", filename="GigachadFour
 ckptgiga = torch.load(modelgiga, map_location=lambda storage, loc: storage)
 generatorgiga.load_state_dict(ckptgiga, strict=False)
 
-modelsketchtrue = hf_hub_download(repo_id="Abhinowww/Capstone", filename="SketchFourHundredTrue.pt")
+modelsketchtrue = hf_hub_download(repo_id="Abhinowww/Capstone", filename="OGSketchFourHundredTrue.pt")
 ckptsketchtrue = torch.load(modelsketchtrue, map_location=lambda storage, loc: storage)
 generatorsketchtrue.load_state_dict(ckptsketchtrue, strict=False)
 
-modelsketchfalse = hf_hub_download(repo_id="Abhinowww/Capstone", filename="SketchFourHundredFalse.pt")
+modelsketchfalse = hf_hub_download(repo_id="Abhinowww/Capstone", filename="OGSketchFourHundredFalse.pt")
 ckptsketchfalse = torch.load(modelsketchfalse, map_location=lambda storage, loc: storage)
 generatorsketchfalse.load_state_dict(ckptsketchfalse, strict=False)
 
@@ -138,10 +148,16 @@ def inference(img, model):
     my_w = projection(aligned_face, "test.pt", device).unsqueeze(0)
     if model == 'Joker':
         with torch.no_grad():
-            my_sample = generatorjoker(my_w, input_is_latent=True)  
+            my_sample = generatorjokerfalse(my_w, input_is_latent=True)  
+    elif model == 'Joker Preserve':
+        with torch.no_grad():
+            my_sample = generatorjokertrue(my_w, input_is_latent=True)
     elif model == 'Voldemort':
         with torch.no_grad():
-            my_sample = generatorvoldemort(my_w, input_is_latent=True)
+            my_sample = generatorvoldemortfalse(my_w, input_is_latent=True)
+    elif model == 'Voldemort Preserve':
+        with torch.no_grad():
+            my_sample = generatorvoldemorttrue(my_w, input_is_latent=True)
     elif model == 'Pushpa':
         with torch.no_grad():
             my_sample = generatorpushpa(my_w, input_is_latent=True)
@@ -180,4 +196,4 @@ description = "Capstone Project. To use it, simply upload your image, or click o
 # css_code='body{background-image:url("https://picsum.photos/seed/picsum/200/300");}'
 # gr.Interface(lambda x:x, "textbox", "textbox", css=css_code).launch(debug=True)
 
-gr.Interface(inference, [gr.inputs.Image(type="pil"),gr.inputs.Dropdown(choices=['Joker', 'Voldemort', 'Pushpa', 'Gigachad', 'Sketch', 'Sketch Preserve'], type="value", default='Joker', label="Model")], gr.outputs.Image(type="pil"),title=title,description=description,allow_flagging=False,allow_screenshot=False).launch()
+gr.Interface(inference, [gr.inputs.Image(type="pil"),gr.inputs.Dropdown(choices=['Joker', 'Joker Preserve', 'Voldemort', 'Voldemort Preserve', 'Pushpa', 'Gigachad', 'Sketch', 'Sketch Preserve'], type="value", default='Joker', label="Model")], gr.outputs.Image(type="pil"),title=title,description=description,allow_flagging=False,allow_screenshot=False).launch()