Add fallback logic to handle PyTorch models and find actual model files - Add manual download fallback when from_pretrained_fastai fails - List repository files to find .pkl or .pt files - Provide clear error message for PyTorch models

app/main.py

@@ -31,7 +31,7 @@ import gradio as gr
 
 # FastAI imports
 from fastai.vision.all import *
-from huggingface_hub import from_pretrained_fastai
+from huggingface_hub import from_pretrained_fastai, hf_hub_download, list_repo_files
 
 from app.config import settings
 
@@ -103,13 +103,94 @@ async def startup_event():
     try:
         model_id = os.getenv("MODEL_ID", "Hammad712/GAN-Colorization-Model")
         logger.info("🔄 Loading FastAI GAN Colorization Model: %s", model_id)
-        learn = from_pretrained_fastai(model_id)
-        logger.info("✅ Model loaded successfully!")
-        model_load_error = None
+        
+        # Try using from_pretrained_fastai first
+        try:
+            learn = from_pretrained_fastai(model_id)
+            logger.info("✅ Model loaded successfully via from_pretrained_fastai!")
+            model_load_error = None
+        except Exception as e1:
+            logger.warning("from_pretrained_fastai failed: %s. Trying manual download...", str(e1))
+            # Fallback: manually download and load the model file
+            hf_token = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_HUB_TOKEN")
+            
+            # List repository files to find the actual model file
+            model_filenames = []
+            model_type = "fastai"
+            
+            try:
+                repo_files = list_repo_files(repo_id=model_id, token=hf_token)
+                logger.info("Repository files: %s", repo_files)
+                pkl_files = [f for f in repo_files if f.endswith('.pkl')]
+                pt_files = [f for f in repo_files if f.endswith('.pt')]
+                
+                if pkl_files:
+                    model_filenames = pkl_files
+                    logger.info("Found .pkl files in repository: %s", pkl_files)
+                    model_type = "fastai"
+                elif pt_files:
+                    model_filenames = pt_files
+                    logger.info("Found .pt files in repository: %s", pt_files)
+                    model_type = "pytorch"
+                else:
+                    model_filenames = ["model.pkl", "export.pkl", "learner.pkl", "model_export.pkl", "generator.pt"]
+                    model_type = "fastai"
+            except Exception as list_err:
+                logger.warning("Could not list repository files: %s. Trying common filenames...", str(list_err))
+                model_filenames = ["model.pkl", "export.pkl", "learner.pkl", "model_export.pkl", "generator.pt"]
+                model_type = "fastai"
+            
+            # Try to download and load the model file
+            cache_dir = os.environ.get("HF_HOME", "/tmp/hf_cache")
+            model_path = None
+            for filename in model_filenames:
+                try:
+                    model_path = hf_hub_download(
+                        repo_id=model_id,
+                        filename=filename,
+                        cache_dir=cache_dir,
+                        token=hf_token
+                    )
+                    logger.info("Found model file: %s", filename)
+                    if filename.endswith('.pt'):
+                        model_type = "pytorch"
+                    elif filename.endswith('.pkl'):
+                        model_type = "fastai"
+                    break
+                except Exception as dl_err:
+                    logger.debug("Failed to download %s: %s", filename, str(dl_err))
+                    continue
+            
+            if model_path and os.path.exists(model_path):
+                if model_type == "pytorch":
+                    error_msg = (
+                        f"Repository '{model_id}' contains a PyTorch model (.pt file), "
+                        f"not a FastAI model. FastAI models must be .pkl files created with FastAI's export. "
+                        f"Please use a FastAI-compatible colorization model, or switch to a different model backend."
+                    )
+                    logger.error(error_msg)
+                    model_load_error = error_msg
+                    raise RuntimeError(error_msg)
+                else:
+                    logger.info("Loading FastAI model from: %s", model_path)
+                    learn = load_learner(model_path)
+                    logger.info("✅ Model loaded successfully from %s", model_path)
+                    model_load_error = None
+            else:
+                error_msg = (
+                    f"Could not find model file in repository '{model_id}'. "
+                    f"Tried: {', '.join(model_filenames)}. "
+                    f"Original error: {str(e1)}"
+                )
+                logger.error(error_msg)
+                model_load_error = error_msg
+                raise RuntimeError(error_msg)
+                
     except Exception as e:
         error_msg = str(e)
+        if not model_load_error:
+            model_load_error = error_msg
         logger.error("❌ Failed to load model: %s", error_msg)
-        model_load_error = error_msg
         # Don't raise - allow health check to work
 
 @app.on_event("shutdown")