ollamadiffuser 1.2.3__py3-none-any.whl → 2.0.0__py3-none-any.whl

ollamadiffuser/core/inference/strategies/sd15_strategy.py

@@ -0,0 +1,134 @@
+"""SD 1.5 inference strategy"""
+
+import logging
+from typing import Optional
+
+import torch
+from PIL import Image
+
+from ..base import InferenceStrategy, SAFETY_DISABLED_KWARGS
+from ...config.settings import ModelConfig
+
+logger = logging.getLogger(__name__)
+
+
+class SD15Strategy(InferenceStrategy):
+    """Strategy for Stable Diffusion 1.5 models"""
+
+    def load(self, model_config: ModelConfig, device: str) -> bool:
+        try:
+            from diffusers import StableDiffusionPipeline
+
+            self.device = device
+            self.model_config = model_config
+
+            load_kwargs = {**SAFETY_DISABLED_KWARGS}
+            dtype = self._get_dtype(device)
+
+            if model_config.variant == "fp16" and device not in ("cpu", "mps"):
+                load_kwargs["torch_dtype"] = dtype
+                load_kwargs["variant"] = "fp16"
+            else:
+                load_kwargs["torch_dtype"] = torch.float32
+
+            self.pipeline = StableDiffusionPipeline.from_pretrained(
+                model_config.path, **load_kwargs
+            )
+            self._move_to_device(device)
+            self._apply_memory_optimizations()
+
+            logger.info(f"SD 1.5 model {model_config.name} loaded on {self.device}")
+            return True
+        except Exception as e:
+            logger.error(f"Failed to load SD 1.5 model: {e}")
+            return False
+
+    def generate(
+        self,
+        prompt: str,
+        negative_prompt: str = "low quality, bad anatomy, worst quality, low resolution",
+        num_inference_steps: Optional[int] = None,
+        guidance_scale: Optional[float] = None,
+        width: int = 512,
+        height: int = 512,
+        seed: Optional[int] = None,
+        image: Optional[Image.Image] = None,
+        mask_image: Optional[Image.Image] = None,
+        strength: float = 0.75,
+        **kwargs,
+    ) -> Image.Image:
+        if not self.pipeline:
+            raise RuntimeError("Model not loaded")
+
+        params = self.model_config.parameters or {}
+        steps = num_inference_steps if num_inference_steps is not None else params.get("num_inference_steps", 50)
+        guidance = guidance_scale if guidance_scale is not None else params.get("guidance_scale", 7.5)
+
+        # Clamp guidance on MPS (known hardware limitation)
+        if self.device == "mps" and guidance > 6.0:
+            logger.info(f"Clamping guidance_scale from {guidance} to 6.0 for MPS stability")
+            guidance = 6.0
+
+        # Warn about non-native resolutions for SD 1.5 (trained on 512x512)
+        if width > 768 or height > 768:
+            logger.warning(
+                f"SD 1.5 was trained on 512x512. Using {width}x{height} may produce artifacts. "
+                "Consider using 512x512 or 768x768 for best results."
+            )
+
+        generator, used_seed = self._make_generator(seed, self.device)
+
+        gen_kwargs = {
+            "prompt": prompt,
+            "negative_prompt": negative_prompt,
+            "num_inference_steps": steps,
+            "guidance_scale": guidance,
+            "width": width,
+            "height": height,
+            "generator": generator,
+        }
+
+        try:
+            # Handle img2img mode
+            if image is not None and mask_image is not None:
+                return self._inpaint(gen_kwargs, image, mask_image, strength)
+            elif image is not None:
+                return self._img2img(gen_kwargs, image, strength)
+
+            logger.info(f"Generating SD 1.5 image: steps={steps}, guidance={guidance}, seed={used_seed}")
+            output = self.pipeline(**gen_kwargs)
+            return output.images[0]
+        except Exception as e:
+            logger.error(f"SD 1.5 generation failed: {e}")
+            return self._create_error_image(str(e), prompt)
+
+    def _img2img(self, gen_kwargs: dict, image: Image.Image, strength: float) -> Image.Image:
+        """Run img2img generation"""
+        from diffusers import StableDiffusionImg2ImgPipeline
+
+        img2img_pipe = StableDiffusionImg2ImgPipeline(**self.pipeline.components)
+        img2img_pipe = img2img_pipe.to(self.device)
+
+        gen_kwargs.pop("width", None)
+        gen_kwargs.pop("height", None)
+        gen_kwargs["image"] = image
+        gen_kwargs["strength"] = strength
+
+        output = img2img_pipe(**gen_kwargs)
+        return output.images[0]
+
+    def _inpaint(self, gen_kwargs: dict, image: Image.Image, mask_image: Image.Image, strength: float) -> Image.Image:
+        """Run inpainting generation"""
+        from diffusers import StableDiffusionInpaintPipeline
+
+        inpaint_pipe = StableDiffusionInpaintPipeline(**self.pipeline.components)
+        inpaint_pipe = inpaint_pipe.to(self.device)
+
+        gen_kwargs.pop("width", None)
+        gen_kwargs.pop("height", None)
+        gen_kwargs["image"] = image
+        gen_kwargs["mask_image"] = mask_image
+        gen_kwargs["strength"] = strength
+
+        output = inpaint_pipe(**gen_kwargs)
+        return output.images[0]

ollamadiffuser/core/inference/strategies/sd3_strategy.py

@@ -0,0 +1,80 @@
+"""SD 3.x inference strategy"""
+
+import logging
+from typing import Optional
+
+import torch
+from PIL import Image
+
+from ..base import InferenceStrategy, SAFETY_DISABLED_KWARGS
+from ...config.settings import ModelConfig
+
+logger = logging.getLogger(__name__)
+
+
+class SD3Strategy(InferenceStrategy):
+    """Strategy for Stable Diffusion 3.x models"""
+
+    def load(self, model_config: ModelConfig, device: str) -> bool:
+        try:
+            from diffusers import StableDiffusion3Pipeline
+
+            self.device = device
+            self.model_config = model_config
+
+            load_kwargs = {**SAFETY_DISABLED_KWARGS}
+            if model_config.variant == "fp16" and device not in ("cpu", "mps"):
+                load_kwargs["torch_dtype"] = torch.float16
+                load_kwargs["variant"] = "fp16"
+            else:
+                load_kwargs["torch_dtype"] = self._get_dtype(device)
+
+            self.pipeline = StableDiffusion3Pipeline.from_pretrained(
+                model_config.path, **load_kwargs
+            )
+            self._move_to_device(device)
+            self._apply_memory_optimizations()
+
+            logger.info(f"SD3 model {model_config.name} loaded on {self.device}")
+            return True
+        except Exception as e:
+            logger.error(f"Failed to load SD3 model: {e}")
+            return False
+
+    def generate(
+        self,
+        prompt: str,
+        negative_prompt: str = "low quality, bad anatomy, worst quality, low resolution",
+        num_inference_steps: Optional[int] = None,
+        guidance_scale: Optional[float] = None,
+        width: int = 1024,
+        height: int = 1024,
+        seed: Optional[int] = None,
+        **kwargs,
+    ) -> Image.Image:
+        if not self.pipeline:
+            raise RuntimeError("Model not loaded")
+
+        params = self.model_config.parameters or {}
+        steps = num_inference_steps if num_inference_steps is not None else params.get("num_inference_steps", 28)
+        guidance = guidance_scale if guidance_scale is not None else params.get("guidance_scale", 3.5)
+
+        generator, used_seed = self._make_generator(seed, self.device)
+
+        gen_kwargs = {
+            "prompt": prompt,
+            "negative_prompt": negative_prompt,
+            "num_inference_steps": steps,
+            "guidance_scale": guidance,
+            "width": width,
+            "height": height,
+            "generator": generator,
+        }
+
+        try:
+            logger.info(f"Generating SD3 image: steps={steps}, guidance={guidance}, seed={used_seed}")
+            output = self.pipeline(**gen_kwargs)
+            return output.images[0]
+        except Exception as e:
+            logger.error(f"SD3 generation failed: {e}")
+            return self._create_error_image(str(e), prompt)

ollamadiffuser/core/inference/strategies/sdxl_strategy.py

@@ -0,0 +1,131 @@
+"""SDXL inference strategy"""
+
+import logging
+from typing import Optional
+
+import torch
+from PIL import Image
+
+from ..base import InferenceStrategy, SAFETY_DISABLED_KWARGS
+from ...config.settings import ModelConfig
+
+logger = logging.getLogger(__name__)
+
+
+class SDXLStrategy(InferenceStrategy):
+    """Strategy for Stable Diffusion XL models"""
+
+    def load(self, model_config: ModelConfig, device: str) -> bool:
+        try:
+            from diffusers import StableDiffusionXLPipeline
+
+            self.device = device
+            self.model_config = model_config
+
+            load_kwargs = {**SAFETY_DISABLED_KWARGS}
+            if model_config.variant == "fp16" and device not in ("cpu", "mps"):
+                load_kwargs["torch_dtype"] = torch.float16
+                load_kwargs["variant"] = "fp16"
+            elif device == "mps":
+                load_kwargs["torch_dtype"] = torch.float32
+            else:
+                load_kwargs["torch_dtype"] = self._get_dtype(device)
+
+            self.pipeline = StableDiffusionXLPipeline.from_pretrained(
+                model_config.path, **load_kwargs
+            )
+
+            # Apply scheduler override if specified in parameters
+            params = model_config.parameters or {}
+            scheduler_class_name = params.get("scheduler_class")
+            if scheduler_class_name:
+                import diffusers
+                scheduler_cls = getattr(diffusers, scheduler_class_name, None)
+                if scheduler_cls is None:
+                    logger.warning(f"Scheduler '{scheduler_class_name}' not found in diffusers, using default")
+                else:
+                    scheduler_kwargs = params.get("scheduler_kwargs", {})
+                    self.pipeline.scheduler = scheduler_cls.from_config(
+                        self.pipeline.scheduler.config, **scheduler_kwargs
+                    )
+                    logger.info(f"Applied scheduler override: {scheduler_class_name}")
+
+            self._move_to_device(device)
+            self._apply_memory_optimizations()
+
+            logger.info(f"SDXL model {model_config.name} loaded on {self.device}")
+            return True
+        except Exception as e:
+            logger.error(f"Failed to load SDXL model: {e}")
+            return False
+
+    def generate(
+        self,
+        prompt: str,
+        negative_prompt: str = "low quality, bad anatomy, worst quality, low resolution",
+        num_inference_steps: Optional[int] = None,
+        guidance_scale: Optional[float] = None,
+        width: int = 1024,
+        height: int = 1024,
+        seed: Optional[int] = None,
+        image: Optional[Image.Image] = None,
+        mask_image: Optional[Image.Image] = None,
+        strength: float = 0.75,
+        **kwargs,
+    ) -> Image.Image:
+        if not self.pipeline:
+            raise RuntimeError("Model not loaded")
+
+        params = self.model_config.parameters or {}
+        steps = num_inference_steps if num_inference_steps is not None else params.get("num_inference_steps", 50)
+        guidance = guidance_scale if guidance_scale is not None else params.get("guidance_scale", 7.5)
+
+        generator, used_seed = self._make_generator(seed, self.device)
+
+        gen_kwargs = {
+            "prompt": prompt,
+            "negative_prompt": negative_prompt,
+            "num_inference_steps": steps,
+            "guidance_scale": guidance,
+            "width": width,
+            "height": height,
+            "generator": generator,
+        }
+
+        try:
+            if image is not None and mask_image is not None:
+                return self._inpaint(gen_kwargs, image, mask_image, strength)
+            elif image is not None:
+                return self._img2img(gen_kwargs, image, strength)
+
+            logger.info(f"Generating SDXL image: steps={steps}, guidance={guidance}, seed={used_seed}")
+            output = self.pipeline(**gen_kwargs)
+            return output.images[0]
+        except Exception as e:
+            logger.error(f"SDXL generation failed: {e}")
+            return self._create_error_image(str(e), prompt)
+
+    def _img2img(self, gen_kwargs: dict, image: Image.Image, strength: float) -> Image.Image:
+        from diffusers import StableDiffusionXLImg2ImgPipeline
+
+        pipe = StableDiffusionXLImg2ImgPipeline(**self.pipeline.components)
+        pipe = pipe.to(self.device)
+        gen_kwargs.pop("width", None)
+        gen_kwargs.pop("height", None)
+        gen_kwargs["image"] = image
+        gen_kwargs["strength"] = strength
+        output = pipe(**gen_kwargs)
+        return output.images[0]
+
+    def _inpaint(self, gen_kwargs: dict, image: Image.Image, mask_image: Image.Image, strength: float) -> Image.Image:
+        from diffusers import StableDiffusionXLInpaintPipeline
+
+        pipe = StableDiffusionXLInpaintPipeline(**self.pipeline.components)
+        pipe = pipe.to(self.device)
+        gen_kwargs.pop("width", None)
+        gen_kwargs.pop("height", None)
+        gen_kwargs["image"] = image
+        gen_kwargs["mask_image"] = mask_image
+        gen_kwargs["strength"] = strength
+        output = pipe(**gen_kwargs)
+        return output.images[0]

ollamadiffuser/core/inference/strategies/video_strategy.py

@@ -0,0 +1,108 @@
+"""Video (AnimateDiff) inference strategy"""
+
+import logging
+from typing import Optional
+
+import torch
+from PIL import Image
+
+from ..base import InferenceStrategy, SAFETY_DISABLED_KWARGS
+from ...config.settings import ModelConfig
+
+logger = logging.getLogger(__name__)
+
+
+class VideoStrategy(InferenceStrategy):
+    """Strategy for AnimateDiff video models"""
+
+    def load(self, model_config: ModelConfig, device: str) -> bool:
+        try:
+            from diffusers import AnimateDiffPipeline, MotionAdapter, DDIMScheduler
+
+            self.device = device
+            self.model_config = model_config
+
+            load_kwargs = {**SAFETY_DISABLED_KWARGS}
+            dtype = torch.float16 if device != "cpu" else torch.float32
+            load_kwargs["torch_dtype"] = dtype
+
+            # Load motion adapter
+            adapter_path = getattr(model_config, "motion_adapter_path", None)
+            if not adapter_path:
+                adapter_path = "guoyww/animatediff-motion-adapter-v1-5-2"
+            motion_adapter = MotionAdapter.from_pretrained(adapter_path, torch_dtype=dtype)
+            load_kwargs["motion_adapter"] = motion_adapter
+
+            self.pipeline = AnimateDiffPipeline.from_pretrained(
+                model_config.path, **load_kwargs
+            )
+
+            # Configure DDIM scheduler for AnimateDiff
+            self.pipeline.scheduler = DDIMScheduler.from_config(
+                self.pipeline.scheduler.config,
+                clip_sample=False,
+                timestep_spacing="linspace",
+                beta_schedule="linear",
+                steps_offset=1,
+            )
+
+            if device in ("cuda", "mps") and hasattr(self.pipeline, "enable_model_cpu_offload"):
+                # CPU offloading manages device placement itself — don't call _move_to_device
+                self.pipeline.enable_model_cpu_offload(device=device)
+            else:
+                self._move_to_device(device)
+
+            if hasattr(self.pipeline, "enable_vae_slicing"):
+                self.pipeline.enable_vae_slicing()
+
+            logger.info(f"Video model {model_config.name} loaded on {self.device}")
+            return True
+        except Exception as e:
+            logger.error(f"Failed to load video model: {e}")
+            return False
+
+    def generate(
+        self,
+        prompt: str,
+        negative_prompt: str = "low quality, bad anatomy, worst quality, low resolution",
+        num_inference_steps: Optional[int] = None,
+        guidance_scale: Optional[float] = None,
+        width: int = 512,
+        height: int = 512,
+        seed: Optional[int] = None,
+        num_frames: int = 16,
+        **kwargs,
+    ) -> Image.Image:
+        if not self.pipeline:
+            raise RuntimeError("Model not loaded")
+
+        params = self.model_config.parameters or {}
+        steps = min(num_inference_steps if num_inference_steps is not None else params.get("num_inference_steps", 25), 25)
+        guidance = min(
+            guidance_scale if guidance_scale is not None else params.get("guidance_scale", 7.5),
+            7.5,
+        )
+
+        generator, used_seed = self._make_generator(seed, self.device)
+
+        gen_kwargs = {
+            "prompt": prompt,
+            "negative_prompt": negative_prompt,
+            "num_inference_steps": steps,
+            "guidance_scale": guidance,
+            "width": width,
+            "height": height,
+            "num_frames": num_frames,
+            "generator": generator,
+        }
+
+        try:
+            logger.info(f"Generating video: {num_frames} frames, steps={steps}, seed={used_seed}")
+            output = self.pipeline(**gen_kwargs)
+
+            if hasattr(output, "frames") and len(output.frames) > 0:
+                return output.frames[0]
+            return output.images[0]
+        except Exception as e:
+            logger.error(f"Video generation failed: {e}")
+            return self._create_error_image(str(e), prompt)

ollamadiffuser/mcp/server.py

@@ -0,0 +1,184 @@
+"""OllamaDiffuser MCP Server - Model Context Protocol integration."""
+
+import asyncio
+import io
+import logging
+import sys
+from typing import Optional
+
+from ..core.models.manager import model_manager
+
+logger = logging.getLogger(__name__)
+
+
+def _ensure_mcp():
+    """Check that the mcp package is available."""
+    try:
+        import mcp  # noqa: F401
+
+        return True
+    except ImportError:
+        logger.error(
+            "MCP package not installed. Install with: pip install 'ollamadiffuser[mcp]'"
+        )
+        return False
+
+
+def create_mcp_server():
+    """Create and configure the MCP server with all tools."""
+    from mcp.server.fastmcp import FastMCP, Image
+
+    mcp_server = FastMCP(
+        "OllamaDiffuser",
+        instructions=(
+            "Local AI image generation via Stable Diffusion, FLUX, and 30+ models"
+        ),
+    )
+
+    @mcp_server.tool()
+    async def generate_image(
+        prompt: str,
+        model: Optional[str] = None,
+        negative_prompt: str = "low quality, bad anatomy, worst quality, low resolution",
+        width: int = 1024,
+        height: int = 1024,
+        steps: Optional[int] = None,
+        guidance_scale: Optional[float] = None,
+        seed: Optional[int] = None,
+    ) -> Image:
+        """Generate an image from a text prompt using a local diffusion model.
+
+        Args:
+            prompt: Text description of the desired image
+            model: Model to use (auto-loads if needed). Leave empty to use current model.
+            negative_prompt: What to avoid in the image
+            width: Image width in pixels
+            height: Image height in pixels
+            steps: Number of inference steps (model-specific default if omitted)
+            guidance_scale: Guidance scale (model-specific default if omitted)
+            seed: Random seed for reproducibility
+        """
+        if model and model_manager.get_current_model() != model:
+            if not model_manager.is_model_installed(model):
+                raise ValueError(
+                    f"Model '{model}' is not installed. "
+                    f"Install it first: ollamadiffuser pull {model}"
+                )
+            logger.info(f"Loading model: {model}")
+            success = await asyncio.to_thread(model_manager.load_model, model)
+            if not success:
+                raise RuntimeError(f"Failed to load model '{model}'")
+
+        if not model_manager.is_model_loaded():
+            raise RuntimeError(
+                "No model loaded. Load one with: load_model('model-name') "
+                "or pass model= parameter"
+            )
+
+        engine = model_manager.loaded_model
+        result = await asyncio.to_thread(
+            engine.generate_image,
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            num_inference_steps=steps,
+            guidance_scale=guidance_scale,
+            width=width,
+            height=height,
+            seed=seed,
+        )
+
+        buf = io.BytesIO()
+        result.save(buf, format="PNG")
+        return Image(data=buf.getvalue(), format="png")
+
+    @mcp_server.tool()
+    async def list_models() -> str:
+        """List all available and installed image generation models.
+
+        Returns a formatted list showing which models are available to download
+        and which are already installed locally.
+        """
+        available = model_manager.list_available_models()
+        installed = model_manager.list_installed_models()
+        current = model_manager.get_current_model()
+
+        lines = ["Available models:"]
+        for name in sorted(available):
+            status_parts = []
+            if name in installed:
+                status_parts.append("installed")
+            if name == current:
+                status_parts.append("loaded")
+            suffix = f" ({', '.join(status_parts)})" if status_parts else ""
+            lines.append(f"  - {name}{suffix}")
+
+        lines.append(f"\nInstalled: {len(installed)}/{len(available)}")
+        if current:
+            lines.append(f"Currently loaded: {current}")
+
+        return "\n".join(lines)
+
+    @mcp_server.tool()
+    async def load_model(model_name: str) -> str:
+        """Load a specific image generation model into memory.
+
+        Args:
+            model_name: Name of the model to load (must be installed first)
+        """
+        if not model_manager.is_model_installed(model_name):
+            installed = model_manager.list_installed_models()
+            return (
+                f"Model '{model_name}' is not installed. "
+                f"Installed models: {', '.join(installed) if installed else 'none'}. "
+                f"Use 'ollamadiffuser pull {model_name}' to install it first."
+            )
+
+        success = await asyncio.to_thread(model_manager.load_model, model_name)
+        if success:
+            return f"Model '{model_name}' loaded successfully"
+        return f"Failed to load model '{model_name}'"
+
+    @mcp_server.tool()
+    async def get_status() -> str:
+        """Get the current status of OllamaDiffuser.
+
+        Returns device info, loaded model, and installed model count.
+        """
+        is_loaded = model_manager.is_model_loaded()
+        current = model_manager.get_current_model()
+        installed = model_manager.list_installed_models()
+
+        lines = ["OllamaDiffuser Status:"]
+        lines.append(f"  Model loaded: {'yes' if is_loaded else 'no'}")
+        if current:
+            lines.append(f"  Current model: {current}")
+        lines.append(f"  Installed models: {len(installed)}")
+
+        if is_loaded and model_manager.loaded_model:
+            engine = model_manager.loaded_model
+            info = engine.get_model_info()
+            if info:
+                lines.append(f"  Device: {info.get('device', 'unknown')}")
+                lines.append(f"  Model type: {info.get('type', 'unknown')}")
+
+        return "\n".join(lines)
+
+    return mcp_server
+
+
+def main():
+    """Entry point for the MCP server (stdio transport)."""
+    if not _ensure_mcp():
+        sys.exit(1)
+
+    logging.basicConfig(
+        level=logging.INFO,
+        format="%(asctime)s - %(levelname)s - %(message)s",
+    )
+
+    server = create_mcp_server()
+    server.run(transport="stdio")
+
+
+if __name__ == "__main__":
+    main()