List repository files to find actual model filename instead of guessing

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, hf_hub_download
+from huggingface_hub import from_pretrained_fastai, hf_hub_download, list_repo_files
 
 from app.config import settings
 
@@ -64,21 +64,37 @@ class ColorizeModel:
             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
+                # First, list files in the repository to find the actual model file
+                hf_token = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_HUB_TOKEN")
+                try:
+                    repo_files = list_repo_files(repo_id=self.model_id, token=hf_token)
+                    logger.info("Repository files: %s", repo_files)
+                    # Look for .pkl files
+                    pkl_files = [f for f in repo_files if f.endswith('.pkl')]
+                    if not pkl_files:
+                        # Also try common FastAI model file names
+                        model_filenames = ["model.pkl", "export.pkl", "learner.pkl", "model_export.pkl"]
+                    else:
+                        model_filenames = pkl_files
+                        logger.info("Found .pkl files in repository: %s", pkl_files)
+                except Exception as list_err:
+                    logger.warning("Could not list repository files: %s. Trying common filenames...", str(list_err))
+                    # Fallback to common filenames
+                    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")
+                            token=hf_token
                         )
                         logger.info("Found model file: %s", filename)
                         break
-                    except Exception:
+                    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):
@@ -87,7 +103,7 @@ class ColorizeModel:
                     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
+                    # If no model file found, raise error with more details
                     raise RuntimeError(
                         f"Could not find model file in repository '{self.model_id}'. "
                         f"Tried: {', '.join(model_filenames)}. "