Detect PyTorch models and provide clear error message - Repository contains generator.pt (PyTorch) not FastAI model

app/colorize_model.py

@@ -69,18 +69,27 @@ class ColorizeModel:
                 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
+                    # Look for .pkl files (FastAI) or .pt files (PyTorch)
                     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:
+                    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:
+                        # Fallback to common filenames
+                        model_filenames = ["model.pkl", "export.pkl", "learner.pkl", "model_export.pkl", "generator.pt"]
+                        model_type = "fastai"  # Default assumption
                 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_filenames = ["model.pkl", "export.pkl", "learner.pkl", "model_export.pkl", "generator.pt"]
+                    model_type = "fastai"
                 
                 model_path = None
                 for filename in model_filenames:
@@ -92,16 +101,33 @@ class ColorizeModel:
                             token=hf_token
                         )
                         logger.info("Found model file: %s", filename)
+                        # Determine model type from extension
+                        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):
-                    # 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)
+                    if model_type == "pytorch":
+                        # Load PyTorch model - this is a GAN generator
+                        logger.info("Loading PyTorch model from: %s", model_path)
+                        # Note: This requires knowing the model architecture
+                        # For now, we'll try to load it and see if it works
+                        logger.warning("PyTorch model loading not fully implemented. This model may not work correctly.")
+                        raise RuntimeError(
+                            f"Repository '{self.model_id}' contains a PyTorch model (generator.pt), "
+                            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."
+                        )
+                    else:
+                        # Load the model using FastAI's load_learner
+                        logger.info("Loading FastAI 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, raise error with more details
                     raise RuntimeError(