Improved Hugging Face Logs

UtkarshTheDev · UtkarshTheDev · commit 4e7e57cf75f5 · 2025-05-08T18:59:27.000+05:30
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -2,14 +2,19 @@
 
 All notable changes to LocalLab will be documented in this file.
 
-## [0.6.2] - 2024-05-03
+## [0.6.2] - 2024-05-04
+
+### Improved
+
+- Improved model downloading experience by using HuggingFace's native progress bars
+- Fixed interleaved progress bars issue during model downloads
+- Added clear success messages after model downloads
 
 ### Fixed
 
 - Fixed CLI configuration issue where optimization settings shown as enabled by default weren't being properly saved
-- Updated default values for all optimization settings (quantization, flash attention, CPU offloading, better transformer) to be enabled by default
+- Updated default values for all optimization settings to be enabled by default
 - Ensured consistency between displayed optimization settings and saved configuration
-- Fixed resource check functions to use correct default values for optimization settings
 
 ## [0.6.1] - 2024-05-02
 
diff --git a/locallab/__init__.py b/locallab/__init__.py
@@ -2,7 +2,7 @@
 LocalLab - A lightweight AI inference server for running LLMs locally
 """
 
-__version__ = "0.6.2"  # Updated to fix CLI optimization settings defaults
+__version__ = "0.6.2"  # Updated to improve model downloading experience and fix CLI settings
 
 # Only import what's necessary initially, lazy-load the rest
 from .logger import get_logger
diff --git a/locallab/logger/__init__.py b/locallab/logger/__init__.py
@@ -109,13 +109,37 @@ class SubduedColoredFormatter(logging.Formatter):
     """Formatter that adds subdued colors to regular logs and bright colors to important logs"""
 
     def format(self, record):
+        # Check if this is a HuggingFace Hub progress bar log
+        # HuggingFace progress bars use tqdm which writes directly to stdout/stderr
+        # We need to completely bypass our logger for these messages
+
         # Check if we're currently downloading a model
         try:
             from ..utils.progress import is_model_downloading
-            if is_model_downloading():
-                # During model download, only show critical logs
-                if record.levelno < logging.ERROR:
-                    # Skip non-critical logs during model download
+
+            # Check if this is a HuggingFace progress bar log
+            is_hf_progress_log = False
+            if hasattr(record, 'name') and isinstance(record.name, str):
+                # HuggingFace Hub logs typically come from these modules
+                hf_modules = ['huggingface_hub', 'filelock', 'transformers', 'tqdm']
+                is_hf_progress_log = any(module in record.name for module in hf_modules)
+
+            # If we're downloading a model and this is a HuggingFace log, skip our formatting
+            if is_model_downloading() and is_hf_progress_log:
+                # Return empty string to skip this log in our logger
+                # HuggingFace will handle displaying its own progress bars
+                return ""
+
+            # For non-HuggingFace logs during model download, only show critical and model-related logs
+            elif is_model_downloading() and record.levelno < logging.ERROR:
+                # Check if this is a model-related log that should be shown
+                is_model_log = False
+                if hasattr(record, 'msg') and isinstance(record.msg, str):
+                    model_patterns = ['model', 'download', 'tokenizer', 'weight']
+                    is_model_log = any(pattern in record.msg.lower() for pattern in model_patterns)
+
+                # Skip non-critical and non-model logs during model download
+                if not is_model_log:
                     return ""
         except (ImportError, AttributeError):
             # If we can't import the function, continue as normal
diff --git a/locallab/model_manager.py b/locallab/model_manager.py
@@ -264,7 +264,19 @@ async def _load_model_with_optimizations(self, model_id: str):
             # Log model loading start
             logger.info(f"Starting download and loading of model: {model_id}")
             print(f"\n{Fore.GREEN}Downloading model: {model_id}{Style.RESET_ALL}")
-            print(f"{Fore.CYAN}This may take a while depending on your internet speed...{Style.RESET_ALL}\n")
+            print(f"{Fore.CYAN}This may take a while depending on your internet speed...{Style.RESET_ALL}")
+            # Add an empty line to separate from HuggingFace progress bars
+            print("")
+
+            # Set a flag to indicate we're downloading a model
+            # This will help our logger know to let HuggingFace's progress bars through
+            try:
+                # Access the module's global variable
+                import locallab.utils.progress
+                locallab.utils.progress.is_downloading = True
+            except:
+                # Fallback if import fails
+                pass
 
             # Load tokenizer first
             logger.info(f"Loading tokenizer for {model_id}...")
@@ -281,7 +293,19 @@ async def _load_model_with_optimizations(self, model_id: str):
                 token=hf_token if hf_token else None,
                 **quant_config
             )
-            print(f"\n{Fore.GREEN}Model {model_id} downloaded successfully!{Style.RESET_ALL}")
+            # Reset the downloading flag
+            try:
+                # Access the module's global variable
+                import locallab.utils.progress
+                locallab.utils.progress.is_downloading = False
+            except:
+                # Fallback if import fails
+                pass
+
+            # Add an empty line and a clear success message
+            print(f"\n{Fore.GREEN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━{Style.RESET_ALL}")
+            print(f"{Fore.GREEN}✅ Model {model_id} downloaded successfully!{Style.RESET_ALL}")
+            print(f"{Fore.GREEN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━{Style.RESET_ALL}")
             logger.info(f"Model weights loaded successfully")
 
             # Apply additional optimizations
@@ -1045,7 +1069,18 @@ async def load_custom_model(self, model_name: str, fallback_model: Optional[str]
             # Load tokenizer first
             logger.info(f"Loading tokenizer for custom model {model_name}...")
             print(f"\n{Fore.GREEN}Downloading custom model: {model_name}{Style.RESET_ALL}")
-            print(f"{Fore.CYAN}This may take a while depending on your internet speed...{Style.RESET_ALL}\n")
+            print(f"{Fore.CYAN}This may take a while depending on your internet speed...{Style.RESET_ALL}")
+            # Add an empty line to separate from HuggingFace progress bars
+            print("")
+
+            # Set a flag to indicate we're downloading a model
+            try:
+                # Access the module's global variable
+                import locallab.utils.progress
+                locallab.utils.progress.is_downloading = True
+            except:
+                # Fallback if import fails
+                pass
 
             self.tokenizer = AutoTokenizer.from_pretrained(model_name)
             logger.info(f"Tokenizer loaded successfully")
@@ -1058,7 +1093,19 @@ async def load_custom_model(self, model_name: str, fallback_model: Optional[str]
                 device_map="auto",
                 quantization_config=quant_config
             )
-            print(f"\n{Fore.GREEN}Custom model {model_name} downloaded successfully!{Style.RESET_ALL}")
+            # Reset the downloading flag
+            try:
+                # Access the module's global variable
+                import locallab.utils.progress
+                locallab.utils.progress.is_downloading = False
+            except:
+                # Fallback if import fails
+                pass
+
+            # Add an empty line and a clear success message
+            print(f"\n{Fore.GREEN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━{Style.RESET_ALL}")
+            print(f"{Fore.GREEN}✅ Custom model {model_name} downloaded successfully!{Style.RESET_ALL}")
+            print(f"{Fore.GREEN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━{Style.RESET_ALL}")
             logger.info(f"Model weights loaded successfully")
 
             self.model = self._apply_optimizations(self.model)
diff --git a/locallab/utils/progress.py b/locallab/utils/progress.py
@@ -155,32 +155,61 @@ def custom_progress_callback(
         pbar.close()
 
 def configure_hf_hub_progress():
-    """Configure HuggingFace Hub to use our custom progress callback"""
+    """
+    Configure HuggingFace Hub to use its native progress bars for model downloads.
+    This completely bypasses our custom logger for HuggingFace download progress.
+    """
     try:
-        from huggingface_hub import constants
+        # 1. Enable HuggingFace's native progress bars
+        from huggingface_hub.utils import logging as hf_logging
+        hf_logging.enable_progress_bars()
 
-        # Disable HF Transfer (which uses a different progress reporting mechanism)
-        constants.HF_HUB_ENABLE_HF_TRANSFER = False
+        # 2. Enable HF Transfer for better download experience
+        from huggingface_hub import constants
+        constants.HF_HUB_ENABLE_HF_TRANSFER = True
 
-        # Configure the download progress callback
+        # 3. Make sure we're NOT overriding HuggingFace's progress callback
+        # This is critical - we want to use their native implementation
         from huggingface_hub import file_download
-        file_download._tqdm_callback = custom_progress_callback
+        if hasattr(file_download, "_tqdm_callback") and file_download._tqdm_callback == custom_progress_callback:
+            # Reset to default if we previously set it to our custom callback
+            file_download._tqdm_callback = None
 
-        # Disable default progress bars
-        try:
-            # Try to disable the default tqdm in huggingface_hub
-            from huggingface_hub.utils import logging as hf_logging
-            hf_logging.disable_progress_bars()
-        except:
-            pass
+        # 4. Set a flag to indicate we're using HuggingFace's native progress bars
+        global is_downloading
+        is_downloading = True
 
-        logger.debug("Configured HuggingFace Hub to use custom progress callback")
+        logger.debug("Configured HuggingFace Hub to use its native progress bars")
     except ImportError:
-        logger.warning("Failed to configure HuggingFace Hub progress callback")
+        logger.warning("Failed to configure HuggingFace Hub progress bars")
     except Exception as e:
-        logger.warning(f"Error configuring HuggingFace Hub progress callback: {str(e)}")
+        logger.warning(f"Error configuring HuggingFace Hub progress: {str(e)}")
 
 # Function to check if we're currently downloading
 def is_model_downloading():
-    """Check if a model is currently being downloaded"""
-    return is_downloading
+    """
+    Check if a model is currently being downloaded.
+
+    This function now checks for active HuggingFace downloads by looking
+    for tqdm progress bars in sys.stdout that contain model file patterns.
+    """
+    # First check our global flag
+    if is_downloading:
+        return True
+
+    # Also check if there are any active HuggingFace downloads
+    # by looking for specific patterns in the output
+    try:
+        # Check if there are any tqdm instances in sys.stdout
+        if hasattr(sys.stdout, '_instances') and sys.stdout._instances:
+            for instance in sys.stdout._instances:
+                if hasattr(instance, 'desc') and isinstance(instance.desc, str):
+                    # Look for common model file patterns in the description
+                    if any(pattern in instance.desc.lower() for pattern in
+                          ['model', 'weight', 'safetensors', 'bin', 'pytorch_model']):
+                        return True
+    except:
+        # If anything goes wrong with the check, default to False
+        pass
+
+    return False
diff --git a/setup.py b/setup.py
@@ -47,7 +47,7 @@
 
 setup(
     name="locallab",
-    version="0.6.2",
+    version="0.6.3",
     packages=find_packages(include=["locallab", "locallab.*"]),
     install_requires=install_requires,
     extras_require={
