ollamadiffuser 1.0.0__py3-none-any.whl → 1.1.1__py3-none-any.whl

ollamadiffuser/core/utils/controlnet_preprocessors.py

@@ -0,0 +1,317 @@
+import cv2
+import numpy as np
+from PIL import Image
+import logging
+from typing import Optional, Union, Tuple, Dict, Any
+import torch
+
+logger = logging.getLogger(__name__)
+
+class ControlNetPreprocessorManager:
+    """Lazy-loading manager for ControlNet preprocessors"""
+    
+    def __init__(self):
+        self.processors = {}
+        self._initialized = False
+        self._initialization_attempted = False
+        self._available_types = []
+    
+    def is_initialized(self) -> bool:
+        """Check if preprocessors are initialized"""
+        return self._initialized
+    
+    def is_available(self) -> bool:
+        """Check if ControlNet preprocessors are available"""
+        if not self._initialization_attempted:
+            # Try a lightweight check without full initialization
+            try:
+                import controlnet_aux
+                return True
+            except ImportError:
+                return False
+        return self._initialized
+    
+    def initialize(self, force: bool = False) -> bool:
+        """
+        Initialize ControlNet preprocessors
+        
+        Args:
+            force: Force re-initialization even if already initialized
+            
+        Returns:
+            True if initialization successful, False otherwise
+        """
+        if self._initialized and not force:
+            return True
+            
+        if self._initialization_attempted and not force:
+            return self._initialized
+            
+        self._initialization_attempted = True
+        logger.info("Initializing ControlNet preprocessors...")
+        
+        try:
+            from controlnet_aux import (
+                CannyDetector,
+                MidasDetector,
+                OpenposeDetector,
+                HEDdetector,
+                MLSDdetector,
+                NormalBaeDetector,
+                LineartDetector,
+                LineartAnimeDetector,
+                ContentShuffleDetector,
+                ZoeDetector
+            )
+            
+            # Initialize processors with proper error handling
+            self.processors = {}
+            
+            # Canny detector (no model download needed)
+            try:
+                self.processors['canny'] = CannyDetector()
+                logger.info("Canny detector initialized")
+            except Exception as e:
+                logger.warning(f"Failed to initialize Canny detector: {e}")
+            
+            # Depth detectors with fallback
+            try:
+                # Try MiDaS first with default model
+                self.processors['depth'] = MidasDetector.from_pretrained('lllyasviel/Annotators')
+                logger.info("MiDaS depth detector initialized")
+            except Exception as e:
+                logger.warning(f"Failed to initialize MiDaS detector: {e}")
+                try:
+                    # Try ZoeDepth as fallback
+                    self.processors['depth_zoe'] = ZoeDetector.from_pretrained('lllyasviel/Annotators')
+                    self.processors['depth'] = self.processors['depth_zoe']  # Use as main depth
+                    logger.info("ZoeDepth detector initialized as fallback")
+                except Exception as e2:
+                    logger.warning(f"Failed to initialize ZoeDepth detector: {e2}")
+            
+            # OpenPose detector
+            try:
+                self.processors['openpose'] = OpenposeDetector.from_pretrained('lllyasviel/Annotators')
+                logger.info("OpenPose detector initialized")
+            except Exception as e:
+                logger.warning(f"Failed to initialize OpenPose detector: {e}")
+            
+            # HED detector
+            try:
+                self.processors['hed'] = HEDdetector.from_pretrained('lllyasviel/Annotators')
+                logger.info("HED detector initialized")
+            except Exception as e:
+                logger.warning(f"Failed to initialize HED detector: {e}")
+            
+            # MLSD detector
+            try:
+                self.processors['mlsd'] = MLSDdetector.from_pretrained('lllyasviel/Annotators')
+                logger.info("MLSD detector initialized")
+            except Exception as e:
+                logger.warning(f"Failed to initialize MLSD detector: {e}")
+            
+            # Normal detector
+            try:
+                self.processors['normal'] = NormalBaeDetector.from_pretrained('lllyasviel/Annotators')
+                logger.info("Normal detector initialized")
+            except Exception as e:
+                logger.warning(f"Failed to initialize Normal detector: {e}")
+            
+            # Lineart detectors
+            try:
+                self.processors['lineart'] = LineartDetector.from_pretrained('lllyasviel/Annotators')
+                logger.info("Lineart detector initialized")
+            except Exception as e:
+                logger.warning(f"Failed to initialize Lineart detector: {e}")
+            
+            try:
+                self.processors['lineart_anime'] = LineartAnimeDetector.from_pretrained('lllyasviel/Annotators')
+                logger.info("Lineart Anime detector initialized")
+            except Exception as e:
+                logger.warning(f"Failed to initialize Lineart Anime detector: {e}")
+            
+            # Content shuffle (no model download needed)
+            try:
+                self.processors['shuffle'] = ContentShuffleDetector()
+                logger.info("Content Shuffle detector initialized")
+            except Exception as e:
+                logger.warning(f"Failed to initialize Content Shuffle detector: {e}")
+            
+            # Add scribble as alias for HED
+            if 'hed' in self.processors:
+                self.processors['scribble'] = self.processors['hed']
+            
+            if self.processors:
+                self._initialized = True
+                self._available_types = list(self.processors.keys())
+                logger.info(f"ControlNet preprocessors initialized: {self._available_types}")
+                return True
+            else:
+                logger.warning("No ControlNet preprocessors could be initialized, falling back to basic processors")
+                self._init_basic_processors()
+                return True
+            
+        except ImportError as e:
+            logger.warning(f"controlnet-aux not available: {e}")
+            # Fallback to basic OpenCV-based processors
+            self._init_basic_processors()
+            return True
+        except Exception as e:
+            logger.error(f"Error initializing ControlNet preprocessors: {e}")
+            # Fallback to basic OpenCV-based processors
+            self._init_basic_processors()
+            return True
+    
+    def _init_basic_processors(self):
+        """Initialize basic OpenCV-based processors as fallback"""
+        logger.info("Using basic OpenCV-based preprocessors")
+        self.processors = {
+            'canny': self._canny_opencv,
+            'depth': self._depth_basic,
+            'scribble': self._scribble_basic,
+        }
+        self._initialized = True
+        self._available_types = list(self.processors.keys())
+    
+    def _canny_opencv(self, image: Image.Image, low_threshold: int = 100, high_threshold: int = 200) -> Image.Image:
+        """Basic Canny edge detection using OpenCV"""
+        # Convert PIL to OpenCV format
+        image_cv = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
+        
+        # Apply Canny edge detection
+        edges = cv2.Canny(image_cv, low_threshold, high_threshold)
+        
+        # Convert back to PIL
+        edges_rgb = cv2.cvtColor(edges, cv2.COLOR_GRAY2RGB)
+        return Image.fromarray(edges_rgb)
+    
+    def _depth_basic(self, image: Image.Image) -> Image.Image:
+        """Basic depth estimation using simple gradients"""
+        # Convert to grayscale
+        gray = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2GRAY)
+        
+        # Apply Gaussian blur
+        blurred = cv2.GaussianBlur(gray, (5, 5), 0)
+        
+        # Create a simple depth map using gradients
+        grad_x = cv2.Sobel(blurred, cv2.CV_64F, 1, 0, ksize=3)
+        grad_y = cv2.Sobel(blurred, cv2.CV_64F, 0, 1, ksize=3)
+        
+        # Combine gradients
+        depth = np.sqrt(grad_x**2 + grad_y**2)
+        depth = np.uint8(255 * depth / np.max(depth))
+        
+        # Convert to RGB
+        depth_rgb = cv2.cvtColor(depth, cv2.COLOR_GRAY2RGB)
+        return Image.fromarray(depth_rgb)
+    
+    def _scribble_basic(self, image: Image.Image) -> Image.Image:
+        """Basic scribble detection using edge detection"""
+        # Convert to grayscale
+        gray = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2GRAY)
+        
+        # Apply edge detection
+        edges = cv2.Canny(gray, 50, 150)
+        
+        # Dilate to make lines thicker
+        kernel = np.ones((3, 3), np.uint8)
+        edges = cv2.dilate(edges, kernel, iterations=1)
+        
+        # Convert to RGB
+        edges_rgb = cv2.cvtColor(edges, cv2.COLOR_GRAY2RGB)
+        return Image.fromarray(edges_rgb)
+    
+    def preprocess(self, 
+                   image: Union[Image.Image, str], 
+                   control_type: str,
+                   **kwargs) -> Image.Image:
+        """
+        Preprocess image for ControlNet
+        
+        Args:
+            image: Input image (PIL Image or path)
+            control_type: Type of control (canny, depth, openpose, etc.)
+            **kwargs: Additional parameters for specific processors
+            
+        Returns:
+            Preprocessed control image
+        """
+        # Initialize if not already done
+        if not self._initialized:
+            if not self.initialize():
+                raise RuntimeError("Failed to initialize ControlNet preprocessors")
+        
+        # Load image if path is provided
+        if isinstance(image, str):
+            image = Image.open(image).convert('RGB')
+        elif not isinstance(image, Image.Image):
+            raise ValueError("Image must be PIL Image or file path")
+        
+        # Ensure image is RGB
+        if image.mode != 'RGB':
+            image = image.convert('RGB')
+        
+        # Get processor
+        if control_type not in self.processors:
+            available = list(self.processors.keys())
+            raise ValueError(f"Control type '{control_type}' not available. Available: {available}")
+        
+        processor = self.processors[control_type]
+        
+        try:
+            if callable(processor):
+                # Basic OpenCV processor
+                return processor(image, **kwargs)
+            else:
+                # ControlNet-aux processor
+                return processor(image, **kwargs)
+                
+        except Exception as e:
+            logger.error(f"Failed to preprocess image with {control_type}: {e}")
+            # Return original image as fallback
+            return image
+    
+    def get_available_types(self) -> list:
+        """Get list of available control types"""
+        if not self._initialized:
+            # Return cached types if available, otherwise return basic types
+            if self._available_types:
+                return self._available_types
+            elif self.is_available():
+                return ['canny', 'depth', 'openpose', 'hed', 'mlsd', 'normal', 'lineart', 'lineart_anime', 'shuffle', 'scribble']
+            else:
+                return ['canny', 'depth', 'scribble']  # Basic OpenCV types
+        return list(self.processors.keys())
+    
+    def resize_for_controlnet(self, 
+                             image: Image.Image, 
+                             width: int = 512, 
+                             height: int = 512) -> Image.Image:
+        """Resize image for ControlNet while maintaining aspect ratio"""
+        # Calculate aspect ratio
+        aspect_ratio = image.width / image.height
+        
+        if aspect_ratio > 1:
+            # Landscape
+            new_width = width
+            new_height = int(width / aspect_ratio)
+        else:
+            # Portrait
+            new_height = height
+            new_width = int(height * aspect_ratio)
+        
+        # Resize image
+        resized = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
+        
+        # Create new image with target size and paste resized image
+        result = Image.new('RGB', (width, height), (0, 0, 0))
+        
+        # Calculate position to center the image
+        x = (width - new_width) // 2
+        y = (height - new_height) // 2
+        
+        result.paste(resized, (x, y))
+        return result
+
+# Global manager instance - no initialization at import time
+controlnet_preprocessor = ControlNetPreprocessorManager() 

ollamadiffuser/core/utils/download_utils.py

@@ -11,25 +11,107 @@ from pathlib import Path
 from huggingface_hub import snapshot_download, hf_hub_download, HfApi
 from tqdm import tqdm
 import threading
+import requests
 
 logger = logging.getLogger(__name__)
 
-class ProgressTracker:
-    """Track download progress across multiple files"""
+class EnhancedProgressTracker:
+    """Enhanced progress tracker that provides Ollama-style detailed progress information"""
     
     def __init__(self, total_files: int = 0, progress_callback: Optional[Callable] = None):
         self.total_files = total_files
         self.completed_files = 0
         self.current_file = ""
         self.file_progress = {}
+        self.file_start_times = {}
+        self.file_speeds = {}
         self.progress_callback = progress_callback
         self.lock = threading.Lock()
+        self.overall_start_time = time.time()
+        self.total_size = 0
+        self.downloaded_size = 0
         
+    def set_total_size(self, total_size: int):
+        """Set the total size for all files"""
+        with self.lock:
+            self.total_size = total_size
+    
+    def start_file(self, filename: str, file_size: int = 0):
+        """Mark a file as started"""
+        with self.lock:
+            self.current_file = filename
+            self.file_start_times[filename] = time.time()
+            self.file_progress[filename] = (0, file_size)
+            
+            # Extract hash-like identifier for Ollama-style display
+            import re
+            hash_match = re.search(r'([a-f0-9]{8,})', filename)
+            if hash_match:
+                display_name = hash_match.group(1)[:12]  # First 12 characters
+            else:
+                # Fallback to filename without extension
+                display_name = Path(filename).stem[:12]
+            
+            if self.progress_callback:
+                self.progress_callback(f"pulling {display_name}")
+    
     def update_file_progress(self, filename: str, downloaded: int, total: int):
-        """Update progress for a specific file"""
+        """Update progress for a specific file with speed calculation"""
         with self.lock:
+            current_time = time.time()
+            
+            # Update file progress
+            old_downloaded = self.file_progress.get(filename, (0, 0))[0]
             self.file_progress[filename] = (downloaded, total)
-            self._report_progress()
+            
+            # Update overall downloaded size
+            size_diff = downloaded - old_downloaded
+            self.downloaded_size += size_diff
+            
+            # Calculate speed for this file
+            if filename in self.file_start_times:
+                elapsed = current_time - self.file_start_times[filename]
+                if elapsed > 0 and downloaded > 0:
+                    speed = downloaded / elapsed  # bytes per second
+                    self.file_speeds[filename] = speed
+            
+            # Report progress in Ollama style
+            if self.progress_callback and total > 0:
+                percentage = (downloaded / total) * 100
+                
+                # Format sizes
+                downloaded_mb = downloaded / (1024 * 1024)
+                total_mb = total / (1024 * 1024)
+                
+                # Calculate speed in MB/s
+                speed_mbps = self.file_speeds.get(filename, 0) / (1024 * 1024)
+                
+                # Calculate ETA
+                if speed_mbps > 0:
+                    remaining_mb = total_mb - downloaded_mb
+                    eta_seconds = remaining_mb / speed_mbps
+                    eta_min = int(eta_seconds // 60)
+                    eta_sec = int(eta_seconds % 60)
+                    eta_str = f"{eta_min}m{eta_sec:02d}s"
+                else:
+                    eta_str = "?"
+                
+                # Extract hash for display
+                import re
+                hash_match = re.search(r'([a-f0-9]{8,})', filename)
+                if hash_match:
+                    display_name = hash_match.group(1)[:12]
+                else:
+                    display_name = Path(filename).stem[:12]
+                
+                # Create progress bar
+                bar_width = 20
+                filled = int((percentage / 100) * bar_width)
+                bar = "█" * filled + " " * (bar_width - filled)
+                
+                progress_msg = f"pulling {display_name}: {percentage:3.0f}% ▕{bar}▏ {downloaded_mb:.0f} MB/{total_mb:.0f} MB {speed_mbps:.0f} MB/s {eta_str}"
+                
+                self.progress_callback(progress_msg)
     
     def complete_file(self, filename: str):
         """Mark a file as completed"""
@@ -38,34 +120,42 @@ class ProgressTracker:
             if filename in self.file_progress:
                 downloaded, total = self.file_progress[filename]
                 self.file_progress[filename] = (total, total)
-            self._report_progress()
-    
-    def set_current_file(self, filename: str):
-        """Set the currently downloading file"""
-        with self.lock:
-            self.current_file = filename
-            self._report_progress()
+            
+            # Report completion
+            if self.progress_callback:
+                import re
+                hash_match = re.search(r'([a-f0-9]{8,})', filename)
+                if hash_match:
+                    display_name = hash_match.group(1)[:12]
+                else:
+                    display_name = Path(filename).stem[:12]
+                
+                total_mb = self.file_progress.get(filename, (0, 0))[1] / (1024 * 1024)
+                self.progress_callback(f"pulling {display_name}: 100% ▕████████████████████▏ {total_mb:.0f} MB/{total_mb:.0f} MB")
     
-    def _report_progress(self):
-        """Report current progress"""
+    def report_overall_progress(self):
+        """Report overall progress"""
         if self.progress_callback:
-            # Calculate overall progress
-            total_downloaded = 0
-            total_size = 0
-            
-            for downloaded, size in self.file_progress.values():
-                total_downloaded += downloaded
-                total_size += size
-            
-            progress_msg = f"Files: {self.completed_files}/{self.total_files}"
-            if total_size > 0:
-                percent = (total_downloaded / total_size) * 100
-                progress_msg += f" | Overall: {percent:.1f}%"
-            
-            if self.current_file:
-                progress_msg += f" | Current: {self.current_file}"
-            
-            self.progress_callback(progress_msg)
+            if self.total_size > 0:
+                overall_percent = (self.downloaded_size / self.total_size) * 100
+                downloaded_gb = self.downloaded_size / (1024 * 1024 * 1024)
+                total_gb = self.total_size / (1024 * 1024 * 1024)
+                
+                elapsed = time.time() - self.overall_start_time
+                if elapsed > 0:
+                    overall_speed = self.downloaded_size / elapsed / (1024 * 1024)  # MB/s
+                    
+                    if overall_speed > 0:
+                        remaining_gb = total_gb - downloaded_gb
+                        eta_seconds = (remaining_gb * 1024) / overall_speed  # Convert GB to MB for calculation
+                        eta_min = int(eta_seconds // 60)
+                        eta_sec = int(eta_seconds % 60)
+                        eta_str = f"{eta_min}m{eta_sec:02d}s"
+                    else:
+                        eta_str = "?"
+                    
+                    progress_msg = f"Overall progress: {overall_percent:.1f}% | {downloaded_gb:.1f} GB/{total_gb:.1f} GB | {overall_speed:.1f} MB/s | ETA: {eta_str}"
+                    self.progress_callback(progress_msg)
 
 def configure_hf_environment():
     """Configure HuggingFace Hub environment for better downloads"""
@@ -132,7 +222,7 @@ def robust_snapshot_download(
     
     # Get file list and sizes for progress tracking
     if progress_callback:
-        progress_callback("📋 Getting repository information...")
+        progress_callback("pulling manifest")
     
     file_sizes = get_repo_file_list(repo_id)
     total_size = sum(file_sizes.values())
@@ -140,20 +230,65 @@ def robust_snapshot_download(
     if progress_callback and file_sizes:
         progress_callback(f"📦 Repository: {len(file_sizes)} files, {format_size(total_size)} total")
     
+    # Initialize enhanced progress tracker
+    progress_tracker = EnhancedProgressTracker(len(file_sizes), progress_callback)
+    progress_tracker.set_total_size(total_size)
+    
     # Check what's already downloaded
     local_path = Path(local_dir)
+    existing_size = 0
     if local_path.exists() and not force_download:
         existing_files = []
-        existing_size = 0
         for file_path in local_path.rglob('*'):
             if file_path.is_file():
                 rel_path = file_path.relative_to(local_path)
                 existing_files.append(str(rel_path))
-                existing_size += file_path.stat().st_size
+                file_size = file_path.stat().st_size
+                existing_size += file_size
+                # Mark existing files as completed in progress tracker
+                progress_tracker.file_progress[str(rel_path)] = (file_size, file_size)
+                progress_tracker.downloaded_size += file_size
+                progress_tracker.completed_files += 1
         
         if progress_callback and existing_files:
             progress_callback(f"📁 Found {len(existing_files)} existing files ({format_size(existing_size)})")
     
+    # Custom tqdm class to capture HuggingFace download progress
+    class OllamaStyleTqdm(tqdm):
+        def __init__(self, *args, **kwargs):
+            # Extract description to get filename
+            desc = kwargs.get('desc', '')
+            self.current_filename = desc
+            
+            # Get file size from our pre-fetched data
+            file_size = file_sizes.get(self.current_filename, 0)
+            if file_size > 0:
+                kwargs['total'] = file_size
+                
+            super().__init__(*args, **kwargs)
+            
+            # Start tracking this file
+            if self.current_filename and progress_callback:
+                progress_tracker.start_file(self.current_filename, file_size)
+        
+        def update(self, n=1):
+            super().update(n)
+            
+            # Update our progress tracker
+            if self.current_filename and progress_callback:
+                downloaded = getattr(self, 'n', 0)
+                total = getattr(self, 'total', 0) or file_sizes.get(self.current_filename, 0)
+                
+                if total > 0:
+                    progress_tracker.update_file_progress(self.current_filename, downloaded, total)
+        
+        def close(self):
+            super().close()
+            
+            # Mark file as completed
+            if self.current_filename and progress_callback:
+                progress_tracker.complete_file(self.current_filename)
+    
     last_exception = None
     
     for attempt in range(max_retries):
@@ -166,12 +301,6 @@ def robust_snapshot_download(
             
             logger.info(f"Download attempt {attempt + 1}/{max_retries} with {workers} workers")
             
-            # Create a custom progress callback for tqdm
-            def tqdm_callback(t):
-                def inner(chunk_size):
-                    t.update(chunk_size)
-                return inner
-            
             result = snapshot_download(
                 repo_id=repo_id,
                 local_dir=local_dir,
@@ -181,7 +310,7 @@ def robust_snapshot_download(
                 resume_download=True,  # Enable resume
                 etag_timeout=300 + (attempt * 60),  # Increase timeout on retries
                 force_download=force_download,
-                tqdm_class=tqdm if progress_callback else None
+                tqdm_class=OllamaStyleTqdm if progress_callback else None
             )
             
             if progress_callback:
@@ -296,8 +425,23 @@ def check_download_integrity(local_dir: str, repo_id: str) -> bool:
         if not local_path.exists():
             return False
         
-        # Check for essential files
-        essential_files = ['model_index.json']
+        # Determine model type based on repo_id
+        is_controlnet = 'controlnet' in repo_id.lower()
+        
+        # Check for essential files based on model type
+        if is_controlnet:
+            # ControlNet models have different essential files
+            essential_files = ['config.json']  # ControlNet models use config.json instead of model_index.json
+            # Also check for model files
+            model_files = ['diffusion_pytorch_model.safetensors', 'diffusion_pytorch_model.bin']
+            has_model_file = any((local_path / model_file).exists() for model_file in model_files)
+            if not has_model_file:
+                logger.warning(f"Missing model file: expected one of {model_files}")
+                return False
+        else:
+            # Regular diffusion models
+            essential_files = ['model_index.json']
+        
         for essential_file in essential_files:
             if not (local_path / essential_file).exists():
                 logger.warning(f"Missing essential file: {essential_file}")
@@ -329,24 +473,29 @@ def check_download_integrity(local_dir: str, repo_id: str) -> bool:
                     logger.warning(f"Empty file detected: {file_path}")
                     return False
         
-        # Check for critical model files
-        critical_dirs = ['transformer', 'text_encoder', 'text_encoder_2', 'tokenizer', 'tokenizer_2']
-        for critical_dir in critical_dirs:
-            dir_path = local_path / critical_dir
-            if dir_path.exists():
-                # Check if directory has any non-empty files
-                has_content = False
-                for file_path in dir_path.rglob('*'):
-                    if file_path.is_file() and file_path.stat().st_size > 0:
-                        # Skip ignored files
-                        should_ignore = any(pattern in str(file_path) for pattern in ignore_patterns)
-                        if not should_ignore:
-                            has_content = True
-                            break
-                
-                if not has_content:
-                    logger.warning(f"Critical directory {critical_dir} appears to be empty or incomplete")
-                    return False
+        # Check for critical model files based on model type
+        if is_controlnet:
+            # ControlNet models are simpler - just need config.json and model weights
+            logger.info("ControlNet model integrity check passed")
+        else:
+            # Check for critical directories in regular diffusion models
+            critical_dirs = ['transformer', 'text_encoder', 'text_encoder_2', 'tokenizer', 'tokenizer_2']
+            for critical_dir in critical_dirs:
+                dir_path = local_path / critical_dir
+                if dir_path.exists():
+                    # Check if directory has any non-empty files
+                    has_content = False
+                    for file_path in dir_path.rglob('*'):
+                        if file_path.is_file() and file_path.stat().st_size > 0:
+                            # Skip ignored files
+                            should_ignore = any(pattern in str(file_path) for pattern in ignore_patterns)
+                            if not should_ignore:
+                                has_content = True
+                                break
+                    
+                    if not has_content:
+                        logger.warning(f"Critical directory {critical_dir} appears to be empty or incomplete")
+                        return False
         
         logger.info("Download integrity check passed")
         return True