Update app.py

app.py

@@ -2,6 +2,7 @@ from huggingface_hub import hf_hub_download
 
 
 Rain_Princess = hf_hub_download(repo_id="maze/FastStyleTransfer", filename="Rain_Princess_512.pth")
+The_Scream = hf_hub_download(repo_id="maze/FastStyleTransfer", filename="Scream_512.pth")
 #modelarcanev3 = hf_hub_download(repo_id="akhaliq/ArcaneGANv0.3", filename="ArcaneGANv0.3.jit")
 #modelarcanev2 = hf_hub_download(repo_id="akhaliq/ArcaneGANv0.2", filename="ArcaneGANv0.2.jit")
 
@@ -136,7 +137,12 @@ def process(image, model):
     
 
 def main(image, backbone, style):
-    transformer.load_state_dict(torch.load(Rain_Princess, map_location=torch.device('cpu')))
+    if style == "The Scream":
+        transformer.load_state_dict(torch.load(The_Scream, map_location=torch.device('cpu')))
+    elif style == "Rain Princess":
+        transformer.load_state_dict(torch.load(Rain_Princess, map_location=torch.device('cpu')))
+    else:
+        transformer.load_state_dict(torch.load(Rain_Princess, map_location=torch.device('cpu')))
     image = Image.fromarray(image)
     isize = image.size
     image = process(image, transformer)