wehub-resource-sync
114 lines
4.1 KiB
Python
114 lines
4.1 KiB
Python
"""
|
|
Demo: MLflow Diffusers Adapter Flavor (LoRA)
|
|
|
|
This script demonstrates the full workflow of logging and loading a diffusion
|
|
model LoRA adapter using the native mlflow.diffusers flavor.
|
|
|
|
No GPU or real model weights required — uses a fake adapter for validation.
|
|
"""
|
|
|
|
import tempfile
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import yaml
|
|
from safetensors.numpy import save_file
|
|
|
|
import mlflow
|
|
import mlflow.diffusers
|
|
|
|
|
|
def create_fake_lora_adapter(output_dir: Path) -> Path:
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
# Simulate LoRA weight matrices (small random tensors)
|
|
tensors = {
|
|
"unet.down_blocks.0.attentions.0.transformer_blocks.0.attn1.to_q.lora_down.weight": (
|
|
np.random.randn(4, 320).astype(np.float32)
|
|
),
|
|
"unet.down_blocks.0.attentions.0.transformer_blocks.0.attn1.to_q.lora_up.weight": (
|
|
np.random.randn(320, 4).astype(np.float32)
|
|
),
|
|
}
|
|
adapter_file = output_dir / "pytorch_lora_weights.safetensors"
|
|
save_file(tensors, str(adapter_file))
|
|
|
|
print(f"Created fake LoRA adapter at: {adapter_file}")
|
|
print(f" Adapter size: {adapter_file.stat().st_size} bytes")
|
|
return output_dir
|
|
|
|
|
|
def demo_log_and_load():
|
|
"""Demonstrate the full log -> load -> inspect cycle."""
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
# 1. Create fake adapter
|
|
adapter_dir = create_fake_lora_adapter(Path(tmpdir) / "my_lora")
|
|
|
|
# 2. Log the adapter with MLflow
|
|
print("\n--- Logging adapter with mlflow.diffusers.log_model() ---")
|
|
mlflow.set_experiment("diffusers-adapter-poc")
|
|
|
|
with mlflow.start_run(run_name="lora-adapter-demo") as run:
|
|
model_info = mlflow.diffusers.log_model(
|
|
adapter_path=str(adapter_dir),
|
|
base_model="black-forest-labs/FLUX.1-dev",
|
|
adapter_type="lora",
|
|
name="lora_model",
|
|
metadata={
|
|
"lora_rank": 4,
|
|
"training_steps": 1000,
|
|
"trigger_word": "sks style",
|
|
},
|
|
)
|
|
|
|
print(f" Run ID: {run.info.run_id}")
|
|
print(f" Model URI: {model_info.model_uri}")
|
|
|
|
# 3. Inspect the MLmodel file
|
|
print("\n--- MLmodel file contents ---")
|
|
model_uri = f"runs:/{run.info.run_id}/lora_model"
|
|
local_path = mlflow.artifacts.download_artifacts(model_uri)
|
|
mlmodel_path = Path(local_path) / "MLmodel"
|
|
|
|
with open(mlmodel_path) as f:
|
|
mlmodel = yaml.safe_load(f)
|
|
|
|
print(yaml.dump(mlmodel, default_flow_style=False, indent=2))
|
|
|
|
# 4. Load the model back
|
|
print("--- Loading model back with mlflow.diffusers.load_model() ---")
|
|
loaded = mlflow.diffusers.load_model(model_uri)
|
|
|
|
print(f" Type: {type(loaded).__name__}")
|
|
print(f" Base model: {loaded.base_model}")
|
|
print(f" Adapter type: {loaded.adapter_type}")
|
|
print(f" Adapter path: {loaded.adapter_path}")
|
|
print(f" Adapter files: {list(Path(loaded.adapter_path).iterdir())}")
|
|
|
|
# 5. Verify flavor config from MLmodel
|
|
print("\n--- Flavor config ---")
|
|
flavor_conf = mlmodel["flavors"]["diffusers"]
|
|
print(f" base_model: {flavor_conf['base_model']}")
|
|
print(f" adapter_type: {flavor_conf['adapter_type']}")
|
|
print(f" adapter_weights: {flavor_conf['adapter_weights']}")
|
|
|
|
# 6. Show that pyfunc interface is available
|
|
print("\n--- Pyfunc model interface ---")
|
|
print(" mlflow.pyfunc.load_model() would return a wrapper with predict()")
|
|
print(" predict() accepts: DataFrame/dict with 'prompt' column")
|
|
print(" predict() returns: list of PNG-encoded image bytes")
|
|
print(" (Skipping actual pyfunc load — requires base model download)")
|
|
|
|
print("\n--- Demo complete! ---")
|
|
print(
|
|
"The adapter is logged as a first-class MLflow model with full model registry support."
|
|
)
|
|
print(
|
|
"To generate images, call loaded.load_pipeline() on a machine "
|
|
"with the base model available."
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
demo_log_and_load()
|