Update app.py

app.py

@@ -4,7 +4,8 @@ import os
 import sys
 import logging
 from pathlib import Path
-from huggingface_hub import snapshot_download, HfHubError
+from huggingface_hub import snapshot_download
+from huggingface_hub.utils import RepositoryNotFoundError, HfHubError
 import torch
 import fire
 import gradio as gr
@@ -69,7 +70,10 @@ def setup_model():
                     token=os.getenv("HF_TOKEN")  # Add token if needed
                 )
                 break
-            except HfHubError as e:
+            except RepositoryNotFoundError as e:
+                logger.error(f"Repository not found: {str(e)}")
+                raise
+            except Exception as e:
                 if attempt == max_retries - 1:
                     logger.error(f"Failed to download model after {max_retries} attempts: {str(e)}")
                     raise