first commit

card_generator_adapter.py

@@ -0,0 +1,107 @@
+"""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