ollamadiffuser 1.1.6__py3-none-any.whl → 1.2.1__py3-none-any.whl

ollamadiffuser/core/models/manager.py

@@ -6,262 +6,28 @@ import logging
 import hashlib
 from huggingface_hub import login
 from ..config.settings import settings, ModelConfig
+from ..config.model_registry import model_registry
 from ..utils.download_utils import robust_snapshot_download, robust_file_download
+from .gguf_loader import gguf_loader, GGUF_AVAILABLE
 
 logger = logging.getLogger(__name__)
 
 class ModelManager:
-    """Model manager"""
+    """Model manager with dynamic registry support and GGUF compatibility"""
     
     def __init__(self):
         self.loaded_model: Optional[object] = None
         self.current_model_name: Optional[str] = None
-        
-        # Predefined model registry
-        self.model_registry = {
-            "flux.1-dev": {
-                "repo_id": "black-forest-labs/FLUX.1-dev",
-                "model_type": "flux",
-                "variant": "bf16",
-                "parameters": {
-                    "num_inference_steps": 16,
-                    "guidance_scale": 2.0,
-                    "max_sequence_length": 512
-                },
-                "hardware_requirements": {
-                    "min_vram_gb": 12,
-                    "recommended_vram_gb": 16,
-                    "min_ram_gb": 24,
-                    "recommended_ram_gb": 32,
-                    "disk_space_gb": 15,
-                    "supported_devices": ["CUDA", "MPS"],
-                    "performance_notes": "Requires NVIDIA RTX 4070+ or Apple M2 Pro+. Needs HuggingFace token. Use 'lora pull' to add LoRA styles."
-                },
-                "license_info": {
-                    "type": "FLUX.1-dev Non-Commercial License",
-                    "requires_agreement": True,
-                    "commercial_use": False
-                }
-            },
-
-            "flux.1-schnell": {
-                "repo_id": "black-forest-labs/FLUX.1-schnell",
-                "model_type": "flux",
-                "variant": "bf16",
-                "parameters": {
-                    "num_inference_steps": 4,
-                    "guidance_scale": 0.0,
-                    "max_sequence_length": 256
-                },
-                "hardware_requirements": {
-                    "min_vram_gb": 12,
-                    "recommended_vram_gb": 16,
-                    "min_ram_gb": 24,
-                    "recommended_ram_gb": 32,
-                    "disk_space_gb": 15,
-                    "supported_devices": ["CUDA", "MPS"],
-                    "performance_notes": "Fast distilled version of FLUX.1-dev. Generates images in ~4 steps. Requires NVIDIA RTX 4070+ or Apple M2 Pro+."
-                },
-                "license_info": {
-                    "type": "Apache 2.0",
-                    "requires_agreement": False,
-                    "commercial_use": True
-                }
-            },
-
-            "stable-diffusion-3.5-medium": {
-                "repo_id": "stabilityai/stable-diffusion-3.5-medium",
-                "model_type": "sd3",
-                "variant": "fp16",
-                "parameters": {
-                    "num_inference_steps": 28,
-                    "guidance_scale": 3.5
-                },
-                "hardware_requirements": {
-                    "min_vram_gb": 8,
-                    "recommended_vram_gb": 12,
-                    "min_ram_gb": 16,
-                    "recommended_ram_gb": 32,
-                    "disk_space_gb": 10,
-                    "supported_devices": ["CUDA", "MPS", "CPU"],
-                    "performance_notes": "Best on NVIDIA RTX 3080+ or Apple M2 Pro+"
-                }
-            },
-            "stable-diffusion-xl-base": {
-                "repo_id": "stabilityai/stable-diffusion-xl-base-1.0",
-                "model_type": "sdxl",
-                "variant": "fp16",
-                "parameters": {
-                    "num_inference_steps": 50,
-                    "guidance_scale": 7.5
-                },
-                "hardware_requirements": {
-                    "min_vram_gb": 6,
-                    "recommended_vram_gb": 10,
-                    "min_ram_gb": 12,
-                    "recommended_ram_gb": 24,
-                    "disk_space_gb": 7,
-                    "supported_devices": ["CUDA", "MPS", "CPU"],
-                    "performance_notes": "Good on NVIDIA RTX 3070+ or Apple M1 Pro+"
-                }
-            },
-            "stable-diffusion-1.5": {
-                "repo_id": "runwayml/stable-diffusion-v1-5",
-                "model_type": "sd15",
-                "variant": "fp16",
-                "parameters": {
-                    "num_inference_steps": 50,
-                    "guidance_scale": 7.5
-                },
-                "hardware_requirements": {
-                    "min_vram_gb": 4,
-                    "recommended_vram_gb": 6,
-                    "min_ram_gb": 8,
-                    "recommended_ram_gb": 16,
-                    "disk_space_gb": 5,
-                    "supported_devices": ["CUDA", "MPS", "CPU"],
-                    "performance_notes": "Runs well on most modern GPUs, including GTX 1060+"
-                }
-            },
-
-            # ControlNet models for SD 1.5
-            "controlnet-canny-sd15": {
-                "repo_id": "lllyasviel/sd-controlnet-canny",
-                "model_type": "controlnet_sd15",
-                "base_model": "stable-diffusion-1.5",
-                "controlnet_type": "canny",
-                "variant": "fp16",
-                "parameters": {
-                    "num_inference_steps": 50,
-                    "guidance_scale": 7.5,
-                    "controlnet_conditioning_scale": 1.0
-                },
-                "hardware_requirements": {
-                    "min_vram_gb": 6,
-                    "recommended_vram_gb": 8,
-                    "min_ram_gb": 12,
-                    "recommended_ram_gb": 20,
-                    "disk_space_gb": 7,
-                    "supported_devices": ["CUDA", "MPS", "CPU"],
-                    "performance_notes": "Requires base SD 1.5 model + ControlNet model. Good for edge detection."
-                }
-            },
-
-            "controlnet-depth-sd15": {
-                "repo_id": "lllyasviel/sd-controlnet-depth",
-                "model_type": "controlnet_sd15",
-                "base_model": "stable-diffusion-1.5",
-                "controlnet_type": "depth",
-                "variant": "fp16",
-                "parameters": {
-                    "num_inference_steps": 50,
-                    "guidance_scale": 7.5,
-                    "controlnet_conditioning_scale": 1.0
-                },
-                "hardware_requirements": {
-                    "min_vram_gb": 6,
-                    "recommended_vram_gb": 8,
-                    "min_ram_gb": 12,
-                    "recommended_ram_gb": 20,
-                    "disk_space_gb": 7,
-                    "supported_devices": ["CUDA", "MPS", "CPU"],
-                    "performance_notes": "Requires base SD 1.5 model + ControlNet model. Good for depth-based control."
-                }
-            },
-
-            "controlnet-openpose-sd15": {
-                "repo_id": "lllyasviel/sd-controlnet-openpose",
-                "model_type": "controlnet_sd15",
-                "base_model": "stable-diffusion-1.5",
-                "controlnet_type": "openpose",
-                "variant": "fp16",
-                "parameters": {
-                    "num_inference_steps": 50,
-                    "guidance_scale": 7.5,
-                    "controlnet_conditioning_scale": 1.0
-                },
-                "hardware_requirements": {
-                    "min_vram_gb": 6,
-                    "recommended_vram_gb": 8,
-                    "min_ram_gb": 12,
-                    "recommended_ram_gb": 20,
-                    "disk_space_gb": 7,
-                    "supported_devices": ["CUDA", "MPS", "CPU"],
-                    "performance_notes": "Requires base SD 1.5 model + ControlNet model. Good for pose control."
-                }
-            },
-
-            "controlnet-scribble-sd15": {
-                "repo_id": "lllyasviel/sd-controlnet-scribble",
-                "model_type": "controlnet_sd15",
-                "base_model": "stable-diffusion-1.5",
-                "controlnet_type": "scribble",
-                "variant": "fp16",
-                "parameters": {
-                    "num_inference_steps": 50,
-                    "guidance_scale": 7.5,
-                    "controlnet_conditioning_scale": 1.0
-                },
-                "hardware_requirements": {
-                    "min_vram_gb": 6,
-                    "recommended_vram_gb": 8,
-                    "min_ram_gb": 12,
-                    "recommended_ram_gb": 20,
-                    "disk_space_gb": 7,
-                    "supported_devices": ["CUDA", "MPS", "CPU"],
-                    "performance_notes": "Requires base SD 1.5 model + ControlNet model. Good for sketch-based control."
-                }
-            },
-
-            # ControlNet models for SDXL
-            "controlnet-canny-sdxl": {
-                "repo_id": "diffusers/controlnet-canny-sdxl-1.0",
-                "model_type": "controlnet_sdxl",
-                "base_model": "stable-diffusion-xl-base",
-                "controlnet_type": "canny",
-                "variant": "fp16",
-                "parameters": {
-                    "num_inference_steps": 50,
-                    "guidance_scale": 7.5,
-                    "controlnet_conditioning_scale": 1.0
-                },
-                "hardware_requirements": {
-                    "min_vram_gb": 8,
-                    "recommended_vram_gb": 12,
-                    "min_ram_gb": 16,
-                    "recommended_ram_gb": 28,
-                    "disk_space_gb": 10,
-                    "supported_devices": ["CUDA", "MPS", "CPU"],
-                    "performance_notes": "Requires base SDXL model + ControlNet model. Good for edge detection with SDXL quality."
-                }
-            },
-
-            "controlnet-depth-sdxl": {
-                "repo_id": "diffusers/controlnet-depth-sdxl-1.0",
-                "model_type": "controlnet_sdxl",
-                "base_model": "stable-diffusion-xl-base",
-                "controlnet_type": "depth",
-                "variant": "fp16",
-                "parameters": {
-                    "num_inference_steps": 50,
-                    "guidance_scale": 7.5,
-                    "controlnet_conditioning_scale": 1.0
-                },
-                "hardware_requirements": {
-                    "min_vram_gb": 8,
-                    "recommended_vram_gb": 12,
-                    "min_ram_gb": 16,
-                    "recommended_ram_gb": 28,
-                    "disk_space_gb": 10,
-                    "supported_devices": ["CUDA", "MPS", "CPU"],
-                    "performance_notes": "Requires base SDXL model + ControlNet model. Good for depth-based control with SDXL quality."
-                }
-            }
-        }
+        self.current_model_type: Optional[str] = None  # Track model type
+    
+    @property
+    def model_registry(self):
+        """Get the current model registry (for backward compatibility)"""
+        return model_registry.get_all_models()
     
     def list_available_models(self) -> List[str]:
         """List all available models"""
-        return list(self.model_registry.keys())
+        return model_registry.get_model_names()
     
     def list_installed_models(self) -> List[str]:
         """List installed models"""
@@ -271,11 +37,24 @@ class ModelManager:
         """Check if model is installed"""
         return model_name in settings.models
     
+    def is_gguf_model(self, model_name: str) -> bool:
+        """Check if a model is a GGUF model"""
+        if not model_name:
+            return False
+        model_info = model_registry.get_model(model_name)
+        if model_info:
+            return gguf_loader.is_gguf_model(model_name, model_info)
+        return False
+    
     def get_model_info(self, model_name: str) -> Optional[Dict]:
         """Get model information"""
-        if model_name in self.model_registry:
-            info = self.model_registry[model_name].copy()
+        info = model_registry.get_model(model_name)
+        if info:
+            # Create a copy to avoid modifying the original
+            info = info.copy()
             info['installed'] = self.is_model_installed(model_name)
+            info['is_gguf'] = self.is_gguf_model(model_name)
+            info['gguf_supported'] = GGUF_AVAILABLE
             if info['installed']:
                 config = settings.models[model_name]
                 info['local_path'] = config.path
@@ -309,13 +88,13 @@ class ModelManager:
                 progress_callback(f"✅ Model {model_name} already installed")
             return True
         
-        if model_name not in self.model_registry:
+        model_info = model_registry.get_model(model_name)
+        if not model_info:
             logger.error(f"Unknown model: {model_name}")
             if progress_callback:
                 progress_callback(f"❌ Error: Unknown model {model_name}")
             return False
         
-        model_info = self.model_registry[model_name]
         model_path = settings.get_model_path(model_name)
         
         # Show model information before download
@@ -373,17 +152,32 @@ class ModelManager:
             if progress_callback:
                 progress_callback(f"🚀 Starting download of {model_name}")
             
+            # Determine download patterns for GGUF models
+            download_kwargs = {
+                "repo_id": model_info["repo_id"],
+                "local_dir": str(model_path),
+                "cache_dir": str(settings.cache_dir),
+                "max_retries": 5,  # Increased retries for large models
+                "initial_workers": 4,  # More workers for faster download
+                "force_download": force,
+                "progress_callback": progress_callback
+            }
+            
+            # Add GGUF-specific file filtering
+            if self.is_gguf_model(model_name):
+                variant = model_info.get("variant", "gguf")
+                patterns = gguf_loader.get_gguf_download_patterns(variant)
+                download_kwargs["allow_patterns"] = patterns["allow_patterns"]
+                download_kwargs["ignore_patterns"] = patterns["ignore_patterns"]
+                
+                if progress_callback:
+                    progress_callback(f"🔍 GGUF model detected - downloading only required files for {variant}")
+                    progress_callback(f"📦 Required files: {len(patterns['allow_patterns'])} files")
+                    progress_callback(f"🚫 Ignoring: {len(patterns['ignore_patterns'])} other GGUF variants")
+            
             # Download main model using robust downloader with enhanced progress
             from ..utils.download_utils import robust_snapshot_download
-            robust_snapshot_download(
-                repo_id=model_info["repo_id"],
-                local_dir=str(model_path),
-                cache_dir=str(settings.cache_dir),
-                max_retries=5,  # Increased retries for large models
-                initial_workers=4,  # More workers for faster download
-                force_download=force,
-                progress_callback=progress_callback
-            )
+            robust_snapshot_download(**download_kwargs)
             
             # Download components (such as LoRA)
             if "components" in model_info:
@@ -397,38 +191,18 @@ class ModelManager:
                     if progress_callback:
                         progress_callback(f"📦 Downloading component: {comp_name}")
                     
-                    if "filename" in comp_info:
-                        # Download single file using robust downloader
-                        from ..utils.download_utils import robust_file_download
-                        robust_file_download(
-                            repo_id=comp_info["repo_id"],
-                            filename=comp_info["filename"],
-                            local_dir=str(comp_path),
-                            cache_dir=str(settings.cache_dir),
-                            max_retries=3,
-                            progress_callback=progress_callback
-                        )
-                    else:
-                        # Download entire repository using robust downloader
-                        robust_snapshot_download(
-                            repo_id=comp_info["repo_id"],
-                            local_dir=str(comp_path),
-                            cache_dir=str(settings.cache_dir),
-                            max_retries=3,
-                            initial_workers=2,  # Use fewer workers for components
-                            force_download=force,
-                            progress_callback=progress_callback
-                        )
-            
-            # Verify download integrity
-            if progress_callback:
-                progress_callback(f"🔍 Verifying download integrity...")
-            
-            from ..utils.download_utils import check_download_integrity
-            if not check_download_integrity(str(model_path), model_info["repo_id"]):
-                raise Exception("Download integrity check failed - some files may be missing or corrupted")
+                    robust_snapshot_download(
+                        repo_id=comp_info["repo_id"],
+                        local_dir=str(comp_path),
+                        cache_dir=str(settings.cache_dir),
+                        allow_patterns=comp_info.get("allow_patterns"),
+                        ignore_patterns=comp_info.get("ignore_patterns", ["*.git*", "README.md", "*.txt"]),
+                        max_retries=3,
+                        initial_workers=2,
+                        progress_callback=progress_callback
+                    )
             
-            # Add to configuration
+            # Create model configuration
             model_config = ModelConfig(
                 name=model_name,
                 path=str(model_path),
@@ -438,26 +212,27 @@ class ModelManager:
                 parameters=model_info.get("parameters")
             )
             
+            # Add to settings
             settings.add_model(model_config)
-            logger.info(f"Model {model_name} download completed")
+            
+            logger.info(f"Model {model_name} downloaded successfully")
             if progress_callback:
-                progress_callback(f"✅ {model_name} download completed successfully and verified!")
+                progress_callback(f"✅ {model_name} downloaded and configured successfully!")
+            
             return True
             
         except Exception as e:
-            logger.error(f"Model download failed: {e}")
+            logger.error(f"Download failed: {str(e)}")
             if progress_callback:
                 progress_callback(f"❌ Download failed: {str(e)}")
             
-            # Clean up failed download only if it's a fresh download
-            if force and model_path.exists():
+            # Clean up partial download
+            if model_path.exists():
                 try:
                     shutil.rmtree(model_path)
-                    logger.info(f"Cleaned up failed download directory: {model_path}")
-                    if progress_callback:
-                        progress_callback(f"🧹 Cleaned up incomplete download")
                 except Exception as cleanup_error:
-                    logger.warning(f"Failed to clean up directory {model_path}: {cleanup_error}")
+                    logger.warning(f"Failed to clean up partial download: {cleanup_error}")
+            
             return False
     
     def remove_model(self, model_name: str) -> bool:
@@ -488,7 +263,7 @@ class ModelManager:
             return False
     
     def load_model(self, model_name: str) -> bool:
-        """Load model into memory"""
+        """Load model into memory (supports both regular and GGUF models)"""
         if not self.is_model_installed(model_name):
             logger.error(f"Model {model_name} is not installed")
             return False
@@ -503,36 +278,72 @@ class ModelManager:
             self.unload_model()
         
         try:
-            from ..inference.engine import InferenceEngine
-            
             model_config = settings.models[model_name]
-            engine = InferenceEngine()
             
-            if engine.load_model(model_config):
-                self.loaded_model = engine
-                self.current_model_name = model_name
-                settings.set_current_model(model_name)
-                logger.info(f"Model {model_name} loaded successfully")
-                return True
+            # Check if this is a GGUF model
+            if self.is_gguf_model(model_name):
+                if not GGUF_AVAILABLE:
+                    logger.error("GGUF support not available. Install with: pip install llama-cpp-python gguf")
+                    return False
+                
+                # Load GGUF model
+                model_config_dict = {
+                    'name': model_name,
+                    'path': model_config.path,
+                    'variant': model_config.variant,
+                    'model_type': model_config.model_type,
+                    'parameters': model_config.parameters
+                }
+                
+                if gguf_loader.load_model(model_config_dict):
+                    self.loaded_model = gguf_loader
+                    self.current_model_name = model_name
+                    self.current_model_type = 'gguf'
+                    settings.set_current_model(model_name)
+                    logger.info(f"GGUF model {model_name} loaded successfully")
+                    return True
+                else:
+                    logger.error(f"GGUF model {model_name} failed to load")
+                    return False
             else:
-                logger.error(f"Model {model_name} failed to load")
-                return False
+                # Load regular diffusion model
+                from ..inference.engine import InferenceEngine
+                
+                engine = InferenceEngine()
+                
+                if engine.load_model(model_config):
+                    self.loaded_model = engine
+                    self.current_model_name = model_name
+                    self.current_model_type = 'diffusion'
+                    settings.set_current_model(model_name)
+                    logger.info(f"Model {model_name} loaded successfully")
+                    return True
+                else:
+                    logger.error(f"Model {model_name} failed to load")
+                    return False
                 
         except Exception as e:
             logger.error(f"Failed to load model: {e}")
             return False
     
     def unload_model(self):
-        """Unload current model"""
+        """Unload current model (supports both regular and GGUF models)"""
         if self.loaded_model is not None:
             try:
-                self.loaded_model.unload()
-                logger.info(f"Model {self.current_model_name} unloaded")
+                if self.current_model_type == 'gguf':
+                    # Unload GGUF model
+                    gguf_loader.unload_model()
+                    logger.info(f"GGUF model {self.current_model_name} unloaded")
+                else:
+                    # Unload regular model
+                    self.loaded_model.unload()
+                    logger.info(f"Model {self.current_model_name} unloaded")
             except Exception as e:
                 logger.error(f"Failed to unload model: {e}")
             finally:
                 self.loaded_model = None
                 self.current_model_name = None
+                self.current_model_type = None
         
         # Also clear the persisted state
         settings.current_model = None
@@ -565,6 +376,22 @@ class ModelManager:
             return response.status_code == 200
         except:
             return False
+    
+    def get_current_model_info(self) -> Optional[Dict]:
+        """Get information about the currently loaded model"""
+        if not self.loaded_model or not self.current_model_name:
+            return None
+            
+        model_info = self.get_model_info(self.current_model_name)
+        if model_info:
+            model_info['loaded'] = True
+            model_info['type'] = self.current_model_type
+            
+            # Add GGUF-specific info if applicable
+            if self.current_model_type == 'gguf':
+                model_info.update(gguf_loader.get_model_info())
+                
+        return model_info
 
 # Global model manager instance
 model_manager = ModelManager()