Update app.py

app.py

@@ -15,11 +15,16 @@ from PIL import Image
 import io
 import time
 
-# Set up a progress callback for huggingface_hub
-from huggingface_hub.utils import RepositoryNotFoundError, RevisionNotFoundError, HfHubHTTPError
-from huggingface_hub.utils._errors import LocalEntryNotFoundError
-from huggingface_hub.file_download import http_get
-from huggingface_hub import HfFileSystem
+# Set up imports for huggingface_hub
+# Import what we can, but handle potential import errors
+try:
+    from huggingface_hub import HfApi, HfFolder, Repository, create_repo, model_info, hf_hub_download
+    from huggingface_hub import login as hf_login
+except ImportError as e:
+    st.error(f"Error importing from huggingface_hub: {e}")
+    
+# Don't rely on specific internal modules that might change between versions
+# We'll handle errors more gracefully instead
 
 # Configure Hugging Face cache and environment
 os.environ["HF_HOME"] = os.path.join(os.getcwd(), ".cache/huggingface")
@@ -138,8 +143,6 @@ def load_model():
     
     # Pre-check model availability
     try:
-        from huggingface_hub import model_info
-        
         # Try to get token first
         token = get_hf_token()
         
@@ -147,9 +150,23 @@ def load_model():
             # Check if the model is available in Hugging Face
             try:
                 model_status.info("Checking model access...")
-                info = model_info("black-forest-labs/FLUX.1-Fill-dev", token=token)
-                model_status.success(f"Model verified: {info.id}")
-                model_status.info(f"Model tags: {', '.join(info.tags)}")
+                
+                # Use safer HfApi instead of direct model_info if possible
+                api = HfApi()
+                try:
+                    info = api.model_info("black-forest-labs/FLUX.1-Fill-dev")
+                    model_status.success(f"Model verified: {info.id}")
+                    if hasattr(info, 'tags'):
+                        model_status.info(f"Model tags: {', '.join(info.tags)}")
+                except Exception as api_error:
+                    # Try fallback with model_info
+                    try:
+                        info = model_info("black-forest-labs/FLUX.1-Fill-dev", token=token)
+                        model_status.success(f"Model verified: {info.id}")
+                        if hasattr(info, 'tags'):
+                            model_status.info(f"Model tags: {', '.join(info.tags)}")
+                    except Exception:
+                        raise api_error
             except Exception as access_error:
                 model_status.error(f"Model access check failed: {str(access_error)}")
         else:
@@ -185,23 +202,34 @@ def load_model():
                     download_progress = st.empty()
                     progress_callback = DownloadProgressCallback(download_progress)
                     
-                    # Set up huggingface_hub to use our progress callback
-                    from huggingface_hub import snapshot_download
-                    
-                    # Download snapshot first to ensure we have all files with progress
+                    # Try pre-downloading files to ensure all components are available
                     progress_placeholder.text("Starting model files download...")
                     try:
-                        # Try to download the model snapshot first
-                        snapshot_download(
-                            repo_id="black-forest-labs/FLUX.1-Fill-dev",
-                            token=token,
-                            local_dir="model_cache",
-                            local_dir_use_symlinks=False,
-                            force_download=False,
-                            resume_download=True,
-                            max_workers=4,
-                            progressbar=True
-                        )
+                        # Try to download individual files one by one instead of snapshot
+                        # This is more robust across different huggingface_hub versions
+                        progress_placeholder.text("Downloading model files using hf_hub_download...")
+                        
+                        # List of essential files to try downloading
+                        essential_files = [
+                            "config.json",
+                            "scheduler/scheduler_config.json",
+                            "tokenizer/tokenizer_config.json"
+                        ]
+                        
+                        # Try downloading files one by one
+                        for file_path in essential_files:
+                            try:
+                                progress_placeholder.text(f"Downloading {file_path}...")
+                                hf_hub_download(
+                                    repo_id="black-forest-labs/FLUX.1-Fill-dev",
+                                    filename=file_path,
+                                    token=token,
+                                    local_dir="model_cache",
+                                    resume_download=True
+                                )
+                            except Exception as file_error:
+                                progress_placeholder.text(f"Error downloading {file_path}: {str(file_error)}")
+                                # Continue with the next file
                         download_progress.text("Download completed successfully!")
                     except Exception as download_error:
                         download_progress.text(f"Snapshot download error: {str(download_error)}")
@@ -323,7 +351,11 @@ def load_model():
                 if token:
                     # First log in explicitly
                     comp_progress.text("Logging in with token...")
-                    login(token)
+                    try:
+                        hf_login(token)
+                    except Exception as login_error:
+                        comp_progress.error(f"Login error (non-critical): {str(login_error)}")
+                        # Continue anyway as the token parameter will be used
                     
                     # Verify login worked and we have access
                     comp_progress.text("Verifying model access...")
@@ -331,8 +363,8 @@ def load_model():
                     try:
                         # Check if we can access the model
                         with timeout(30):  # 30 seconds timeout for API check
-                            model_info = api.model_info("black-forest-labs/FLUX.1-Fill-dev")
-                            comp_progress.text(f"Access verified to model: {model_info.id}")
+                            model_info_obj = api.model_info("black-forest-labs/FLUX.1-Fill-dev")
+                            comp_progress.text(f"Access verified to model: {model_info_obj.id if hasattr(model_info_obj, 'id') else 'black-forest-labs/FLUX.1-Fill-dev'}")
                     except Exception as access_error:
                         st.error(f"Authentication succeeded but access verification failed: {str(access_error)}")
                         if "401" in str(access_error) or "access" in str(access_error).lower():