ollamadiffuser 1.0.0__py3-none-any.whl

ollamadiffuser/cli/main.py

@@ -0,0 +1,597 @@
+#!/usr/bin/env python3
+import click
+import sys
+import logging
+from typing import Optional
+from rich.console import Console
+from rich.table import Table
+from rich.progress import Progress, SpinnerColumn, TextColumn
+from rich import print as rprint
+
+from ..core.models.manager import model_manager
+from ..core.config.settings import settings
+from ..api.server import run_server
+
+console = Console()
+
+@click.group()
+@click.option('--verbose', '-v', is_flag=True, help='Enable verbose output')
+def cli(verbose):
+    """OllamaDiffuser - Image generation model management tool"""
+    if verbose:
+        logging.basicConfig(level=logging.DEBUG)
+    else:
+        logging.basicConfig(level=logging.WARNING)
+
+@cli.command()
+@click.argument('model_name')
+@click.option('--force', '-f', is_flag=True, help='Force re-download')
+def pull(model_name: str, force: bool):
+    """Download model"""
+    rprint(f"[blue]Downloading model: {model_name}[/blue]")
+    
+    with Progress(
+        SpinnerColumn(),
+        TextColumn("[progress.description]{task.description}"),
+        console=console
+    ) as progress:
+        task = progress.add_task(f"Downloading {model_name}...", total=None)
+        
+        def progress_callback(message: str):
+            """Update progress display with download status"""
+            progress.update(task, description=message)
+        
+        if model_manager.pull_model(model_name, force=force, progress_callback=progress_callback):
+            progress.update(task, description=f"✅ {model_name} download completed")
+            rprint(f"[green]Model {model_name} downloaded successfully![/green]")
+        else:
+            progress.update(task, description=f"❌ {model_name} download failed")
+            rprint(f"[red]Model {model_name} download failed![/red]")
+            sys.exit(1)
+
+@cli.command()
+@click.argument('model_name')
+@click.option('--host', '-h', default=None, help='Server host address')
+@click.option('--port', '-p', default=None, type=int, help='Server port')
+def run(model_name: str, host: Optional[str], port: Optional[int]):
+    """Run model service"""
+    rprint(f"[blue]Starting model service: {model_name}[/blue]")
+    
+    # Check if model is installed
+    if not model_manager.is_model_installed(model_name):
+        rprint(f"[red]Model {model_name} is not installed. Please run first: ollamadiffuser pull {model_name}[/red]")
+        sys.exit(1)
+    
+    # Load model
+    rprint("[yellow]Loading model...[/yellow]")
+    if not model_manager.load_model(model_name):
+        rprint(f"[red]Failed to load model {model_name}![/red]")
+        sys.exit(1)
+    
+    rprint(f"[green]Model {model_name} loaded successfully![/green]")
+    
+    # Start server
+    try:
+        run_server(host=host, port=port)
+    except KeyboardInterrupt:
+        rprint("\n[yellow]Server stopped[/yellow]")
+        model_manager.unload_model()
+        # Clear the current model from settings when server stops
+        settings.current_model = None
+        settings.save_config()
+
+@cli.command()
+@click.option('--hardware', '-hw', is_flag=True, help='Show hardware requirements')
+def list(hardware: bool):
+    """List all models"""
+    available_models = model_manager.list_available_models()
+    installed_models = model_manager.list_installed_models()
+    current_model = model_manager.get_current_model()
+    
+    if hardware:
+        # Show detailed hardware requirements
+        for model_name in available_models:
+            info = model_manager.get_model_info(model_name)
+            if not info:
+                continue
+                
+            # Check installation status
+            if model_name in installed_models:
+                status = "✅ Installed"
+                if model_name == current_model:
+                    status += " (current)"
+                size = info.get('size', 'Unknown')
+            else:
+                status = "⬇️ Available"
+                size = "-"
+            
+            # Create individual table for each model
+            table = Table(title=f"[bold cyan]{model_name}[/bold cyan] - {status}")
+            table.add_column("Property", style="yellow", no_wrap=True)
+            table.add_column("Value", style="white")
+            
+            # Basic info
+            table.add_row("Type", info.get('model_type', 'Unknown'))
+            table.add_row("Size", size)
+            
+            # Hardware requirements
+            hw_req = info.get('hardware_requirements', {})
+            if hw_req:
+                table.add_row("Min VRAM", f"{hw_req.get('min_vram_gb', 'Unknown')} GB")
+                table.add_row("Recommended VRAM", f"{hw_req.get('recommended_vram_gb', 'Unknown')} GB")
+                table.add_row("Min RAM", f"{hw_req.get('min_ram_gb', 'Unknown')} GB")
+                table.add_row("Recommended RAM", f"{hw_req.get('recommended_ram_gb', 'Unknown')} GB")
+                table.add_row("Disk Space", f"{hw_req.get('disk_space_gb', 'Unknown')} GB")
+                table.add_row("Supported Devices", ", ".join(hw_req.get('supported_devices', [])))
+                table.add_row("Performance Notes", hw_req.get('performance_notes', 'N/A'))
+            
+            console.print(table)
+            console.print()  # Add spacing between models
+    else:
+        # Show compact table
+        table = Table(title="OllamaDiffuser Model List")
+        table.add_column("Model Name", style="cyan", no_wrap=True)
+        table.add_column("Status", style="green")
+        table.add_column("Size", style="blue")
+        table.add_column("Type", style="magenta")
+        table.add_column("Min VRAM", style="yellow")
+        
+        for model_name in available_models:
+            # Check installation status
+            if model_name in installed_models:
+                status = "✅ Installed"
+                if model_name == current_model:
+                    status += " (current)"
+                
+                # Get model information
+                info = model_manager.get_model_info(model_name)
+                size = info.get('size', 'Unknown') if info else 'Unknown'
+                model_type = info.get('model_type', 'Unknown') if info else 'Unknown'
+            else:
+                status = "⬇️ Available"
+                size = "-"
+                # Get type from registry
+                info = model_manager.get_model_info(model_name)
+                model_type = info.get('model_type', 'Unknown') if info else 'Unknown'
+            
+            # Get hardware requirements
+            hw_req = info.get('hardware_requirements', {}) if info else {}
+            min_vram = f"{hw_req.get('min_vram_gb', '?')} GB" if hw_req else "Unknown"
+            
+            table.add_row(model_name, status, size, model_type, min_vram)
+        
+        console.print(table)
+        console.print("\n[dim]💡 Use --hardware flag to see detailed hardware requirements[/dim]")
+
+@cli.command()
+@click.argument('model_name')
+def show(model_name: str):
+    """Show model detailed information"""
+    info = model_manager.get_model_info(model_name)
+    
+    if info is None:
+        rprint(f"[red]Model {model_name} does not exist[/red]")
+        sys.exit(1)
+    
+    rprint(f"[bold cyan]Model Information: {model_name}[/bold cyan]")
+    rprint(f"Type: {info.get('model_type', 'Unknown')}")
+    rprint(f"Variant: {info.get('variant', 'Unknown')}")
+    rprint(f"Installed: {'Yes' if info.get('installed', False) else 'No'}")
+    
+    if info.get('installed', False):
+        rprint(f"Local Path: {info.get('local_path', 'Unknown')}")
+        rprint(f"Size: {info.get('size', 'Unknown')}")
+    
+    # Hardware requirements
+    if 'hardware_requirements' in info and info['hardware_requirements']:
+        hw_req = info['hardware_requirements']
+        rprint("\n[bold]Hardware Requirements:[/bold]")
+        rprint(f"  Min VRAM: {hw_req.get('min_vram_gb', 'Unknown')} GB")
+        rprint(f"  Recommended VRAM: {hw_req.get('recommended_vram_gb', 'Unknown')} GB")
+        rprint(f"  Min RAM: {hw_req.get('min_ram_gb', 'Unknown')} GB")
+        rprint(f"  Recommended RAM: {hw_req.get('recommended_ram_gb', 'Unknown')} GB")
+        rprint(f"  Disk Space: {hw_req.get('disk_space_gb', 'Unknown')} GB")
+        rprint(f"  Supported Devices: {', '.join(hw_req.get('supported_devices', []))}")
+        if hw_req.get('performance_notes'):
+            rprint(f"  Performance Notes: {hw_req.get('performance_notes')}")
+    
+    if 'parameters' in info and info['parameters']:
+        rprint("\n[bold]Default Parameters:[/bold]")
+        for key, value in info['parameters'].items():
+            rprint(f"  {key}: {value}")
+    
+    if 'components' in info and info['components']:
+        rprint("\n[bold]Components:[/bold]")
+        for key, value in info['components'].items():
+            rprint(f"  {key}: {value}")
+
+@cli.command()
+@click.argument('model_name', required=False)
+@click.option('--list', '-l', is_flag=True, help='List all available models')
+def check(model_name: str, list: bool):
+    """Check model download status and integrity"""
+    if list:
+        rprint("[bold blue]📋 Available Models:[/bold blue]")
+        available_models = model_manager.list_available_models()
+        for model in available_models:
+            model_info = model_manager.get_model_info(model)
+            status = "✅ Installed" if model_manager.is_model_installed(model) else "⬇️ Available"
+            license_type = model_info.get("license_info", {}).get("type", "Unknown")
+            rprint(f"   {model:<30} {status:<15} ({license_type})")
+        return
+    
+    if not model_name:
+        rprint("[bold red]❌ Please specify a model name or use --list[/bold red]")
+        rprint("[dim]Usage: ollamadiffuser check MODEL_NAME[/dim]")
+        rprint("[dim]       ollamadiffuser check --list[/dim]")
+        return
+    
+    # Import and run the check function
+    import subprocess
+    import sys
+    from pathlib import Path
+    
+    # Run the check script
+    script_path = Path(__file__).parent.parent.parent / "examples" / "check_model_download.py"
+    try:
+        result = subprocess.run([sys.executable, str(script_path), model_name], 
+                              capture_output=True, text=True)
+        rprint(result.stdout)
+        if result.stderr:
+            rprint(f"[red]{result.stderr}[/red]")
+    except Exception as e:
+        rprint(f"[red]❌ Error running check: {e}[/red]")
+
+@cli.command()
+@click.argument('model_name')
+@click.confirmation_option(prompt='Are you sure you want to delete this model?')
+def rm(model_name: str):
+    """Remove model"""
+    if model_manager.remove_model(model_name):
+        rprint(f"[green]Model {model_name} removed successfully![/green]")
+    else:
+        rprint(f"[red]Failed to remove model {model_name}![/red]")
+        sys.exit(1)
+
+@cli.command()
+def ps():
+    """Show currently running model"""
+    current_model = model_manager.get_current_model()
+    server_running = model_manager.is_server_running()
+    
+    if current_model:
+        rprint(f"[green]Current model: {current_model}[/green]")
+        
+        # Check server status
+        if server_running:
+            rprint(f"[green]Server status: Running on {settings.server.host}:{settings.server.port}[/green]")
+            
+            # Try to get model info from the running server
+            try:
+                import requests
+                response = requests.get(f"http://{settings.server.host}:{settings.server.port}/api/models/running", timeout=2)
+                if response.status_code == 200:
+                    data = response.json()
+                    if data.get('loaded'):
+                        info = data.get('info', {})
+                        rprint(f"Device: {info.get('device', 'Unknown')}")
+                        rprint(f"Type: {info.get('type', 'Unknown')}")
+                        rprint(f"Variant: {info.get('variant', 'Unknown')}")
+                    else:
+                        rprint("[yellow]Model loaded but not active in server[/yellow]")
+            except:
+                pass
+        else:
+            rprint("[yellow]Server status: Not running[/yellow]")
+            rprint("[dim]Model is set as current but server is not active[/dim]")
+            
+        # Show model info from local config
+        model_info = model_manager.get_model_info(current_model)
+        if model_info:
+            rprint(f"Model type: {model_info.get('model_type', 'Unknown')}")
+            if model_info.get('installed'):
+                rprint(f"Size: {model_info.get('size', 'Unknown')}")
+    else:
+        if server_running:
+            rprint("[yellow]Server is running but no model is loaded[/yellow]")
+            rprint(f"[green]Server status: Running on {settings.server.host}:{settings.server.port}[/green]")
+        else:
+            rprint("[yellow]No model is currently running[/yellow]")
+            rprint("[dim]Use 'ollamadiffuser run <model>' to start a model[/dim]")
+
+@cli.command()
+@click.option('--host', '-h', default=None, help='Server host address')
+@click.option('--port', '-p', default=None, type=int, help='Server port')
+def serve(host: Optional[str], port: Optional[int]):
+    """Start API server (without loading model)"""
+    rprint("[blue]Starting OllamaDiffuser API server...[/blue]")
+    
+    try:
+        run_server(host=host, port=port)
+    except KeyboardInterrupt:
+        rprint("\n[yellow]Server stopped[/yellow]")
+
+@cli.command()
+@click.argument('model_name')
+def load(model_name: str):
+    """Load model into memory"""
+    rprint(f"[blue]Loading model: {model_name}[/blue]")
+    
+    if model_manager.load_model(model_name):
+        rprint(f"[green]Model {model_name} loaded successfully![/green]")
+    else:
+        rprint(f"[red]Failed to load model {model_name}![/red]")
+        sys.exit(1)
+
+@cli.command()
+def unload():
+    """Unload current model"""
+    if model_manager.is_model_loaded():
+        current_model = model_manager.get_current_model()
+        model_manager.unload_model()
+        rprint(f"[green]Model {current_model} unloaded[/green]")
+    else:
+        rprint("[yellow]No model to unload[/yellow]")
+
+@cli.command()
+def stop():
+    """Stop running server"""
+    if not model_manager.is_server_running():
+        rprint("[yellow]No server is currently running[/yellow]")
+        return
+    
+    try:
+        import requests
+        import signal
+        import psutil
+        
+        host = settings.server.host
+        port = settings.server.port
+        
+        # Try graceful shutdown via API first
+        try:
+            response = requests.post(f"http://{host}:{port}/api/shutdown", timeout=5)
+            if response.status_code == 200:
+                rprint("[green]Server stopped gracefully[/green]")
+                return
+        except:
+            pass
+        
+        # Fallback: Find and terminate the process
+        for proc in psutil.process_iter(['pid', 'name', 'cmdline']):
+            try:
+                cmdline = proc.info['cmdline']
+                if cmdline and any('uvicorn' in arg for arg in cmdline) and any(str(port) in arg for arg in cmdline):
+                    proc.terminate()
+                    proc.wait(timeout=10)
+                    rprint(f"[green]Server process (PID: {proc.info['pid']}) stopped[/green]")
+                    return
+            except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.TimeoutExpired):
+                continue
+        
+        rprint("[red]Could not find or stop the server process[/red]")
+        
+    except ImportError:
+        rprint("[red]psutil package required for stop command. Install with: pip install psutil[/red]")
+    except Exception as e:
+        rprint(f"[red]Failed to stop server: {e}[/red]")
+
+@cli.group()
+def lora():
+    """LoRA (Low-Rank Adaptation) management commands"""
+    pass
+
+@lora.command()
+@click.argument('repo_id')
+@click.option('--weight-name', '-w', help='Specific weight file name (e.g., lora.safetensors)')
+@click.option('--alias', '-a', help='Local alias name for the LoRA')
+def pull(repo_id: str, weight_name: Optional[str], alias: Optional[str]):
+    """Download LoRA weights from Hugging Face Hub"""
+    from ..core.utils.lora_manager import lora_manager
+    
+    rprint(f"[blue]Downloading LoRA: {repo_id}[/blue]")
+    
+    with Progress(
+        SpinnerColumn(),
+        TextColumn("[progress.description]{task.description}"),
+        console=console
+    ) as progress:
+        task = progress.add_task(f"Downloading LoRA...", total=None)
+        
+        def progress_callback(message: str):
+            progress.update(task, description=message)
+        
+        if lora_manager.pull_lora(repo_id, weight_name=weight_name, alias=alias, progress_callback=progress_callback):
+            progress.update(task, description=f"✅ LoRA download completed")
+            rprint(f"[green]LoRA {repo_id} downloaded successfully![/green]")
+        else:
+            progress.update(task, description=f"❌ LoRA download failed")
+            rprint(f"[red]LoRA {repo_id} download failed![/red]")
+            sys.exit(1)
+
+@lora.command()
+@click.argument('lora_name')
+@click.option('--scale', '-s', default=1.0, type=float, help='LoRA scale/strength (default: 1.0)')
+def load(lora_name: str, scale: float):
+    """Load LoRA weights into the current model"""
+    from ..core.utils.lora_manager import lora_manager
+    
+    rprint(f"[blue]Loading LoRA: {lora_name} (scale: {scale})[/blue]")
+    
+    if lora_manager.load_lora(lora_name, scale=scale):
+        rprint(f"[green]LoRA {lora_name} loaded successfully![/green]")
+    else:
+        rprint(f"[red]Failed to load LoRA {lora_name}![/red]")
+        sys.exit(1)
+
+@lora.command()
+def unload():
+    """Unload current LoRA weights"""
+    from ..core.utils.lora_manager import lora_manager
+    
+    rprint("[blue]Unloading LoRA weights...[/blue]")
+    
+    if lora_manager.unload_lora():
+        rprint("[green]LoRA weights unloaded successfully![/green]")
+    else:
+        rprint("[red]Failed to unload LoRA weights![/red]")
+        sys.exit(1)
+
+@lora.command()
+@click.argument('lora_name')
+@click.confirmation_option(prompt='Are you sure you want to delete this LoRA?')
+def rm(lora_name: str):
+    """Remove LoRA weights"""
+    from ..core.utils.lora_manager import lora_manager
+    
+    rprint(f"[blue]Removing LoRA: {lora_name}[/blue]")
+    
+    if lora_manager.remove_lora(lora_name):
+        rprint(f"[green]LoRA {lora_name} removed successfully![/green]")
+    else:
+        rprint(f"[red]Failed to remove LoRA {lora_name}![/red]")
+        sys.exit(1)
+
+@lora.command()
+def ps():
+    """Show currently loaded LoRA status"""
+    from ..core.utils.lora_manager import lora_manager
+    
+    # Check if server is running
+    server_running = lora_manager._is_server_running()
+    current_lora = lora_manager.get_current_lora()
+    
+    if server_running:
+        rprint(f"[green]Server status: Running on {settings.server.host}:{settings.server.port}[/green]")
+        
+        # Try to get LoRA status from the running server
+        try:
+            import requests
+            response = requests.get(f"http://{settings.server.host}:{settings.server.port}/api/models/running", timeout=2)
+            if response.status_code == 200:
+                data = response.json()
+                if data.get('loaded'):
+                    model_info = data.get('info', {})
+                    rprint(f"Model: {data.get('model', 'Unknown')}")
+                    rprint(f"Device: {model_info.get('device', 'Unknown')}")
+                    rprint(f"Type: {model_info.get('type', 'Unknown')}")
+                else:
+                    rprint("[yellow]No model loaded in server[/yellow]")
+                    return
+        except Exception as e:
+            rprint(f"[red]Failed to get server status: {e}[/red]")
+            return
+    else:
+        # Check local model manager
+        if model_manager.is_model_loaded():
+            current_model = model_manager.get_current_model()
+            rprint(f"[green]Model loaded locally: {current_model}[/green]")
+        else:
+            rprint("[yellow]No server running and no local model loaded[/yellow]")
+            rprint("[dim]Use 'ollamadiffuser run <model>' to start a model[/dim]")
+            return
+    
+    # Show LoRA status
+    lora_status_shown = False
+    lora_loaded_on_server = False
+    
+    # Try to get LoRA status from server if running
+    if server_running:
+        try:
+            import requests
+            response = requests.get(f"http://{settings.server.host}:{settings.server.port}/api/lora/status", timeout=2)
+            if response.status_code == 200:
+                lora_data = response.json()
+                if lora_data.get('loaded'):
+                    lora_info = lora_data.get('info', {})
+                    rprint(f"\n[bold green]🔄 LoRA Status: LOADED (via server)[/bold green]")
+                    rprint(f"Adapter: {lora_info.get('adapter_name', 'Unknown')}")
+                    if 'scale' in lora_info:
+                        rprint(f"Scale: {lora_info.get('scale', 'Unknown')}")
+                    if 'adapters' in lora_info:
+                        rprint(f"Active Adapters: {', '.join(lora_info.get('adapters', []))}")
+                    lora_status_shown = True
+                    lora_loaded_on_server = True
+                else:
+                    rprint(f"\n[dim]💾 LoRA Status: No LoRA loaded (server)[/dim]")
+                    lora_status_shown = True
+        except Exception as e:
+            rprint(f"\n[yellow]⚠️ Failed to get LoRA status from server: {e}[/yellow]")
+    
+    # Fallback to local LoRA manager state
+    if not lora_status_shown:
+        if current_lora:
+            lora_info = lora_manager.get_lora_info(current_lora)
+            if lora_info:
+                rprint(f"\n[bold green]🔄 LoRA Status: LOADED (local)[/bold green]")
+                rprint(f"Name: {current_lora}")
+                rprint(f"Repository: {lora_info.get('repo_id', 'Unknown')}")
+                rprint(f"Weight File: {lora_info.get('weight_name', 'Unknown')}")
+                rprint(f"Size: {lora_info.get('size', 'Unknown')}")
+                rprint(f"Local Path: {lora_info.get('path', 'Unknown')}")
+            else:
+                rprint(f"\n[yellow]⚠️ LoRA {current_lora} is set as current but info not found[/yellow]")
+        else:
+            rprint(f"\n[dim]💾 LoRA Status: No LoRA loaded[/dim]")
+    
+    if not lora_loaded_on_server:
+        rprint("[dim]Use 'ollamadiffuser lora load <lora_name>' to load a LoRA[/dim]")
+
+@lora.command()
+def list():
+    """List available and installed LoRA weights"""
+    from ..core.utils.lora_manager import lora_manager
+    
+    installed_loras = lora_manager.list_installed_loras()
+    current_lora = lora_manager.get_current_lora()
+    
+    if not installed_loras:
+        rprint("[yellow]No LoRA weights installed.[/yellow]")
+        rprint("\n[dim]💡 Use 'ollamadiffuser lora pull <repo_id>' to download LoRA weights[/dim]")
+        return
+    
+    table = Table(title="Installed LoRA Weights")
+    table.add_column("Name", style="cyan", no_wrap=True)
+    table.add_column("Repository", style="blue")
+    table.add_column("Status", style="green")
+    table.add_column("Size", style="yellow")
+    
+    for lora_name, lora_info in installed_loras.items():
+        status = "🔄 Loaded" if lora_name == current_lora else "💾 Available"
+        size = lora_info.get('size', 'Unknown')
+        repo_id = lora_info.get('repo_id', 'Unknown')
+        
+        table.add_row(lora_name, repo_id, status, size)
+    
+    console.print(table)
+
+@lora.command()
+@click.argument('lora_name')
+def show(lora_name: str):
+    """Show detailed LoRA information"""
+    from ..core.utils.lora_manager import lora_manager
+    
+    lora_info = lora_manager.get_lora_info(lora_name)
+    
+    if not lora_info:
+        rprint(f"[red]LoRA {lora_name} not found.[/red]")
+        sys.exit(1)
+    
+    rprint(f"[bold cyan]LoRA Information: {lora_name}[/bold cyan]")
+    rprint(f"Repository: {lora_info.get('repo_id', 'Unknown')}")
+    rprint(f"Weight File: {lora_info.get('weight_name', 'Unknown')}")
+    rprint(f"Local Path: {lora_info.get('path', 'Unknown')}")
+    rprint(f"Size: {lora_info.get('size', 'Unknown')}")
+    rprint(f"Downloaded: {lora_info.get('downloaded_at', 'Unknown')}")
+    
+    if lora_info.get('description'):
+        rprint(f"Description: {lora_info.get('description')}")
+
+@cli.command()
+def version():
+    """Show version information"""
+    rprint("[bold cyan]OllamaDiffuser v1.0.0[/bold cyan]")
+    rprint("Image generation model management tool")
+
+if __name__ == '__main__':
+    cli()