108 lines
3.6 KiB
Python
108 lines
3.6 KiB
Python
"""Adapter to load the LoRA checkpoint and define conditioning logic.
|
|
|
|
Customize this file to match your model architecture, then use:
|
|
--generator-module card_generator_adapter.py
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any, Mapping
|
|
|
|
|
|
def build_pipeline(checkpoint_path: str, device: str):
|
|
"""Load LoRA adapter and return a callable pipeline.
|
|
|
|
The pipeline must accept:
|
|
pipeline(prompt_or_conditioning, num_inference_steps=30, guidance_scale=7.5)
|
|
|
|
and return an object with .images attribute.
|
|
"""
|
|
from pathlib import Path
|
|
|
|
checkpoint_input = Path(checkpoint_path).expanduser().resolve()
|
|
if checkpoint_input.is_dir():
|
|
checkpoint_dir = checkpoint_input
|
|
elif checkpoint_input.exists():
|
|
checkpoint_dir = checkpoint_input.parent
|
|
else:
|
|
raise FileNotFoundError(f"Checkpoint path not found: {checkpoint_input}")
|
|
|
|
# Load base Stable Diffusion model + LoRA adapter (PEFT)
|
|
try:
|
|
from diffusers import StableDiffusionPipeline
|
|
import torch
|
|
except ImportError as e:
|
|
raise RuntimeError(
|
|
f"diffusers and torch required. Install: pip install diffusers torch "
|
|
f"(error: {e})"
|
|
)
|
|
|
|
# Load base model
|
|
model_id = "runwayml/stable-diffusion-v1-5"
|
|
pipe = StableDiffusionPipeline.from_pretrained(
|
|
model_id,
|
|
torch_dtype=torch.float16 if device == "cuda" else torch.float32,
|
|
)
|
|
pipe = pipe.to(device)
|
|
|
|
# Load LoRA weights from adapter_model.safetensors
|
|
adapter_path = checkpoint_dir / "adapter_model.safetensors"
|
|
if adapter_path.exists():
|
|
try:
|
|
pipe.load_lora_weights(str(checkpoint_dir))
|
|
except Exception as e:
|
|
message = str(e)
|
|
if "PEFT backend is required" in message:
|
|
raise RuntimeError(
|
|
"Failed to load LoRA: PEFT backend is missing. "
|
|
"Install required packages with: pip install peft transformers accelerate safetensors"
|
|
) from e
|
|
raise RuntimeError(
|
|
f"Failed to load LoRA from {checkpoint_dir}: {e}\n"
|
|
"Ensure adapter_config.json and adapter_model.safetensors are present."
|
|
) from e
|
|
else:
|
|
raise FileNotFoundError(
|
|
f"LoRA adapter not found at {adapter_path}. "
|
|
f"Expected: adapter_model.safetensors in {checkpoint_dir}"
|
|
)
|
|
|
|
return pipe
|
|
|
|
|
|
def metadata_to_conditioning(meta: Mapping[str, Any]) -> str:
|
|
"""Convert metadata dict to a Stable Diffusion prompt.
|
|
|
|
LoRA is trained on Pokemon cards, so describe it as such.
|
|
"""
|
|
name = str(meta.get("name", "Unknown Pokemon"))
|
|
pokemon_type = str(meta.get("type", "normal")).capitalize()
|
|
secondary = meta.get("secondary_type")
|
|
|
|
hp = str(meta.get("hp", "60"))
|
|
|
|
attacks = meta.get("attacks") or []
|
|
attack_list = []
|
|
if isinstance(attacks, list):
|
|
for atk in attacks:
|
|
if isinstance(atk, dict):
|
|
attack_list.append(str(atk.get("name", "")).lower())
|
|
elif atk:
|
|
attack_list.append(str(atk).lower())
|
|
|
|
# Build a descriptive prompt for card generation
|
|
prompt = f"Pokemon trading card of {name}, {pokemon_type}-type Pokemon"
|
|
if secondary:
|
|
prompt += f"/{secondary.capitalize()}"
|
|
prompt += f", HP {hp}"
|
|
|
|
if attack_list:
|
|
prompt += f", with attacks: {', '.join(attack_list[:2])}"
|
|
|
|
description = meta.get("description", "").strip()
|
|
if description:
|
|
prompt += f". {description}"
|
|
|
|
prompt += ". High quality illustration, official Pokemon card style."
|
|
return prompt
|