ollamadiffuser 1.0.0__py3-none-any.whl

ollamadiffuser/core/models/manager.py

@@ -0,0 +1,436 @@
+import os
+import shutil
+from pathlib import Path
+from typing import Dict, List, Optional, Union
+import logging
+import hashlib
+from huggingface_hub import login
+from ..config.settings import settings, ModelConfig
+from ..utils.download_utils import robust_snapshot_download, robust_file_download
+
+logger = logging.getLogger(__name__)
+
+class ModelManager:
+    """Model manager"""
+    
+    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": 20,
+                    "guidance_scale": 7.0
+                },
+                "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+"
+                }
+            }
+        }
+    
+    def list_available_models(self) -> List[str]:
+        """List all available models"""
+        return list(self.model_registry.keys())
+    
+    def list_installed_models(self) -> List[str]:
+        """List installed models"""
+        return list(settings.models.keys())
+    
+    def is_model_installed(self, model_name: str) -> bool:
+        """Check if model is installed"""
+        return model_name in settings.models
+    
+    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['installed'] = self.is_model_installed(model_name)
+            if info['installed']:
+                config = settings.models[model_name]
+                info['local_path'] = config.path
+                info['size'] = self._get_model_size(config.path)
+            return info
+        return None
+    
+    def _get_model_size(self, model_path: str) -> str:
+        """Get model size"""
+        try:
+            path = Path(model_path)
+            if path.is_file():
+                size = path.stat().st_size
+            else:
+                size = sum(f.stat().st_size for f in path.rglob('*') if f.is_file())
+            
+            # Convert to human readable format
+            for unit in ['B', 'KB', 'MB', 'GB', 'TB']:
+                if size < 1024.0:
+                    return f"{size:.1f} {unit}"
+                size /= 1024.0
+            return f"{size:.1f} PB"
+        except Exception:
+            return "Unknown"
+    
+    def pull_model(self, model_name: str, force: bool = False, progress_callback=None) -> bool:
+        """Download model using robust download utilities with detailed progress tracking"""
+        if not force and self.is_model_installed(model_name):
+            logger.info(f"Model {model_name} already exists")
+            if progress_callback:
+                progress_callback(f"✅ Model {model_name} already installed")
+            return True
+        
+        if model_name not in self.model_registry:
+            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
+        if progress_callback:
+            license_info = model_info.get("license_info", {})
+            progress_callback(f"📦 Model: {model_name}")
+            progress_callback(f"🔗 Repository: {model_info['repo_id']}")
+            if license_info:
+                progress_callback(f"📄 License: {license_info.get('type', 'Unknown')}")
+                if license_info.get('requires_agreement', False):
+                    progress_callback(f"🔑 HuggingFace token required - ensure HF_TOKEN is set")
+                else:
+                    progress_callback(f"✅ No HuggingFace token required")
+        
+        # Check if partial download exists and is valid
+        if not force and model_path.exists():
+            if progress_callback:
+                progress_callback(f"🔍 Checking existing download...")
+            
+            from ..utils.download_utils import check_download_integrity
+            if check_download_integrity(str(model_path), model_info["repo_id"]):
+                if progress_callback:
+                    progress_callback(f"✅ Found complete download, adding to configuration...")
+                
+                # Add to configuration
+                model_config = ModelConfig(
+                    name=model_name,
+                    path=str(model_path),
+                    model_type=model_info["model_type"],
+                    variant=model_info.get("variant"),
+                    components=model_info.get("components"),
+                    parameters=model_info.get("parameters")
+                )
+                
+                settings.add_model(model_config)
+                logger.info(f"Model {model_name} configuration updated")
+                if progress_callback:
+                    progress_callback(f"✅ {model_name} ready to use!")
+                return True
+            else:
+                if progress_callback:
+                    progress_callback(f"⚠️ Incomplete download detected, will resume...")
+        
+        try:
+            # Ensure HuggingFace token is set
+            if settings.hf_token:
+                login(token=settings.hf_token)
+                if progress_callback:
+                    progress_callback(f"🔑 Authenticated with HuggingFace")
+            else:
+                if progress_callback:
+                    progress_callback(f"⚠️ No HuggingFace token found - some models may not be accessible")
+            
+            logger.info(f"Downloading model: {model_name}")
+            if progress_callback:
+                progress_callback(f"🚀 Starting download of {model_name}")
+            
+            # 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
+            )
+            
+            # Download components (such as LoRA)
+            if "components" in model_info:
+                components_path = model_path / "components"
+                components_path.mkdir(exist_ok=True)
+                
+                for comp_name, comp_info in model_info["components"].items():
+                    comp_path = components_path / comp_name
+                    comp_path.mkdir(exist_ok=True)
+                    
+                    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")
+            
+            # Add to configuration
+            model_config = ModelConfig(
+                name=model_name,
+                path=str(model_path),
+                model_type=model_info["model_type"],
+                variant=model_info.get("variant"),
+                components=model_info.get("components"),
+                parameters=model_info.get("parameters")
+            )
+            
+            settings.add_model(model_config)
+            logger.info(f"Model {model_name} download completed")
+            if progress_callback:
+                progress_callback(f"✅ {model_name} download completed successfully and verified!")
+            return True
+            
+        except Exception as e:
+            logger.error(f"Model download failed: {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():
+                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}")
+            return False
+    
+    def remove_model(self, model_name: str) -> bool:
+        """Remove model"""
+        if not self.is_model_installed(model_name):
+            logger.error(f"Model {model_name} is not installed")
+            return False
+        
+        try:
+            # If currently using this model, unload it first
+            if self.current_model_name == model_name:
+                self.unload_model()
+            
+            # Delete model files
+            model_config = settings.models[model_name]
+            model_path = Path(model_config.path)
+            if model_path.exists():
+                shutil.rmtree(model_path)
+            
+            # Remove from configuration
+            settings.remove_model(model_name)
+            
+            logger.info(f"Model {model_name} has been removed")
+            return True
+            
+        except Exception as e:
+            logger.error(f"Failed to remove model: {e}")
+            return False
+    
+    def load_model(self, model_name: str) -> bool:
+        """Load model into memory"""
+        if not self.is_model_installed(model_name):
+            logger.error(f"Model {model_name} is not installed")
+            return False
+        
+        # If the same model is already loaded, return directly
+        if self.current_model_name == model_name:
+            logger.info(f"Model {model_name} is already loaded")
+            return True
+        
+        # Unload current model
+        if self.loaded_model is not None:
+            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
+            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"""
+        if self.loaded_model is not None:
+            try:
+                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
+        
+        # Also clear the persisted state
+        settings.current_model = None
+        settings.save_config()
+    
+    def get_current_model(self) -> Optional[str]:
+        """Get current loaded model name"""
+        # First check in-memory state
+        if self.current_model_name:
+            return self.current_model_name
+        # Then check persisted state
+        return settings.current_model
+    
+    def is_model_loaded(self) -> bool:
+        """Check if a model is loaded in memory"""
+        # Only check in-memory state - a model is truly loaded only if it's in memory
+        return self.loaded_model is not None
+    
+    def has_current_model(self) -> bool:
+        """Check if there's a current model set (may not be loaded in memory)"""
+        return settings.current_model is not None
+    
+    def is_server_running(self) -> bool:
+        """Check if the server is actually running"""
+        try:
+            import requests
+            host = settings.server.host
+            port = settings.server.port
+            response = requests.get(f"http://{host}:{port}/api/health", timeout=2)
+            return response.status_code == 200
+        except:
+            return False
+
+# Global model manager instance
+model_manager = ModelManager() 

ollamadiffuser/core/utils/__init__.py

@@ -0,0 +1,3 @@
+"""
+Core utilities package
+"""