Add fallback mechanism to manually download and load FastAI model if from_pretrained_fastai fails

app/colorize_model.py

@@ -21,7 +21,7 @@ os.environ["XDG_CACHE_HOME"] = cache_dir
 import torch
 from PIL import Image
 from fastai.vision.all import *
-from huggingface_hub import from_pretrained_fastai
+from huggingface_hub import from_pretrained_fastai, hf_hub_download
 
 from app.config import settings
 
@@ -57,8 +57,42 @@ class ColorizeModel:
 
         logger.info("Loading FastAI GAN Colorization model: %s", self.model_id)
         try:
-            self.learn = from_pretrained_fastai(self.model_id)
-            logger.info("FastAI GAN Colorization model loaded successfully")
+            # Try using from_pretrained_fastai first
+            try:
+                self.learn = from_pretrained_fastai(self.model_id)
+                logger.info("FastAI GAN Colorization model loaded successfully via from_pretrained_fastai")
+            except Exception as e1:
+                logger.warning("from_pretrained_fastai failed: %s. Trying manual download...", str(e1))
+                # Fallback: manually download and load the model file
+                # Try common FastAI model file names
+                model_filenames = ["model.pkl", "export.pkl", "learner.pkl", "model_export.pkl"]
+                model_path = None
+                
+                for filename in model_filenames:
+                    try:
+                        model_path = hf_hub_download(
+                            repo_id=self.model_id,
+                            filename=filename,
+                            cache_dir=self.cache_dir,
+                            token=os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_HUB_TOKEN")
+                        )
+                        logger.info("Found model file: %s", filename)
+                        break
+                    except Exception:
+                        continue
+                
+                if model_path and os.path.exists(model_path):
+                    # Load the model using FastAI's load_learner
+                    logger.info("Loading model from: %s", model_path)
+                    self.learn = load_learner(model_path)
+                    logger.info("FastAI GAN Colorization model loaded successfully from %s", model_path)
+                else:
+                    # If no model file found, try listing repository files
+                    raise RuntimeError(
+                        f"Could not find model file in repository '{self.model_id}'. "
+                        f"Tried: {', '.join(model_filenames)}. "
+                        f"Original error: {str(e1)}"
+                    )
         except Exception as e:
             error_msg = (
                 f"Failed to load FastAI model '{self.model_id}'. "