chore: import upstream snapshot with attribution

mlflow/diffusers/wrapper.py

@@ -0,0 +1,151 @@
+import io
+import logging
+import threading
+from types import MappingProxyType
+from typing import Any
+
+import pandas as pd
+
+from mlflow.diffusers import _detect_device
+from mlflow.exceptions import MlflowException
+
+_logger = logging.getLogger(__name__)
+
+
+class _DiffusersAdapterWrapper:
+    def __init__(
+        self,
+        adapter_path: str,
+        flavor_conf: dict[str, Any],
+        model_config: dict[str, Any] | None = None,
+    ):
+        self._adapter_path = adapter_path
+        self._flavor_conf = flavor_conf
+        self._model_config = MappingProxyType(model_config or {})
+        self._pipeline = None
+        self._load_lock = threading.Lock()
+
+    def _load_pipeline(self):
+        from diffusers import DiffusionPipeline
+
+        base_model = self._model_config.get("base_model") or self._flavor_conf["base_model"]
+        base_model_revision = self._flavor_conf.get("base_model_revision")
+        device = _detect_device(self._model_config.get("device"))
+        torch_dtype = self._model_config.get("torch_dtype", "auto")
+
+        load_kwargs = {"torch_dtype": torch_dtype}
+        if base_model_revision:
+            load_kwargs["revision"] = base_model_revision
+
+        weight_name = self._flavor_conf.get("weight_name")
+        lora_kwargs = {}
+        if weight_name:
+            lora_kwargs["weight_name"] = weight_name
+
+        _logger.info("Loading base pipeline: %s", base_model)
+        try:
+            pipe = DiffusionPipeline.from_pretrained(base_model, **load_kwargs)
+        except OSError as e:
+            raise MlflowException(
+                f"Failed to load base model '{base_model}'. If the model has moved, "
+                "pass the correct location via "
+                "model_config={{'base_model': '<new_path_or_hub_id>'}} "
+                "when loading with mlflow.pyfunc.load_model()."
+            ) from e
+
+        _logger.info("Loading LoRA adapter from: %s", self._adapter_path)
+        pipe.load_lora_weights(self._adapter_path, **lora_kwargs)
+
+        self._pipeline = pipe.to(device)
+
+    def get_raw_model(self):
+        if self._pipeline is None:
+            with self._load_lock:
+                if self._pipeline is None:
+                    self._load_pipeline()
+        return self._pipeline
+
+    def _flatten_prompts(self, prompts):
+        """Flatten nested lists produced by schema enforcement."""
+        flat = []
+        for item in prompts:
+            if isinstance(item, list):
+                flat.extend(item)
+            else:
+                flat.append(item)
+        return flat
+
+    def predict(self, data, params: dict[str, Any] | None = None):
+        pipeline = self.get_raw_model()
+
+        if isinstance(data, pd.DataFrame):
+            if "prompt" in data.columns:
+                prompts = data["prompt"].tolist()
+            elif len(data.columns) == 1:
+                # Schema enforcement wraps scalar strings into a single-column DataFrame
+                prompts = data.iloc[:, 0].tolist()
+            else:
+                raise MlflowException(
+                    f"Input DataFrame must contain a 'prompt' column. "
+                    f"Got columns: {list(data.columns)}"
+                )
+            # Schema enforcement may wrap {"prompt": ["a","b"]} into a
+            # single-row DataFrame where the cell contains a list, producing
+            # [["a","b"]] after tolist(). Flatten to ["a","b"].
+            prompts = self._flatten_prompts(prompts)
+        elif isinstance(data, str):
+            prompts = [data]
+        elif isinstance(data, dict):
+            if "prompt" not in data:
+                raise MlflowException(
+                    f"Input dict must contain a 'prompt' key. Got keys: {list(data.keys())}"
+                )
+            prompts = data["prompt"]
+            if isinstance(prompts, str):
+                prompts = [prompts]
+            elif isinstance(prompts, list):
+                prompts = self._flatten_prompts(prompts)
+            else:
+                raise MlflowException(
+                    "'prompt' value must be a string or list of strings, "
+                    f"got {type(prompts).__name__}."
+                )
+        elif isinstance(data, list):
+            prompts = self._flatten_prompts(data)
+        else:
+            raise MlflowException(f"Unsupported input type: {type(data)}")
+
+        if not prompts:
+            raise MlflowException(
+                "No prompts provided. Input must contain at least one prompt string."
+            )
+
+        if any(p is None for p in prompts):
+            raise MlflowException(
+                "Prompt values must be strings, not None. "
+                "Check your input for missing or null values."
+            )
+
+        params = params or {}
+        param_keys = ("num_inference_steps", "guidance_scale", "height", "width", "negative_prompt")
+        gen_kwargs = {k: params[k] for k in param_keys if k in params}
+        # Drop empty-string negative_prompt so the pipeline uses its own default
+        if gen_kwargs.get("negative_prompt") == "":
+            del gen_kwargs["negative_prompt"]
+
+        output = pipeline(prompt=prompts, **gen_kwargs)
+
+        if not hasattr(output, "images") or not output.images:
+            raise MlflowException(
+                "Pipeline returned no images. The output may have been filtered "
+                "by the safety checker, or the pipeline does not support image generation."
+            )
+
+        results = []
+        for image in output.images:
+            buf = io.BytesIO()
+            image.save(buf, format="PNG")
+            results.append(buf.getvalue())
+            buf.close()
+
+        return results