Update app.py

app.py

@@ -25,37 +25,21 @@ COLOUR_MODEL = "RGB"
 
 MODEL_REPO = "NDugar/horse_to_zebra_cycle_GAN"
 MODEL_FILE = "h2z-85epoch.pth"
-
-# Model Initalisation
-#shinkai_model_hfhub = hf_hub_download(repo_id=MODEL_REPO_SHINKAI, filename=MODEL_FILE_SHINKAI)
-#hosoda_model_hfhub = hf_hub_download(repo_id=MODEL_REPO_HOSODA, filename=MODEL_FILE_HOSODA)
-#miyazaki_model_hfhub = hf_hub_download(repo_id=MODEL_REPO_MIYAZAKI, filename=MODEL_FILE_MIYAZAKI)
 model_hfhub = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE)
-
-#shinkai_model = Transformer()
-#hosoda_model = Transformer()
-#miyazaki_model = Transformer()
-#model = Transformer()
-
 enable_gpu = torch.cuda.is_available()
 map_location = torch.device("cuda") if enable_gpu else "cpu"
+from huggingface_hub import hf_hub_download
+from fastai.learner import load_learner
 
-model.load_state_dict(torch.load(model_hfhub, map_location=map_location))
-
+model = load_learner(
+    hf_hub_download("NDugar/horse_to_zebra_cycle_GAN", "model.pkl")
+)
 model.eval()
-
-
-# Functions
-
 def get_model():
     return model
-
-
 def adjust_image_for_model(img):
     logger.info(f"Image Height: {img.height}, Image Width: {img.width}")
     return img
-
-
 def inference(img, style):
     img = adjust_image_for_model(img)
     input_image = img.convert(COLOUR_MODEL)
@@ -103,4 +87,6 @@ gr.Interface(
     examples=examples,
     allow_flagging="never",
     allow_screenshot=False,
-).launch(enable_queue=True)
+).launch(enable_queue=True)
+
+