AIGC 摄影训练工程化：从打标数据到可控出图的完整管线

import os
from PIL import Image
from transformers import BlipProcessor, BlipForConditionalGeneration

processor = BlipProcessor.from_pretrained(
    "Salesforce/blip-image-captioning-base")
model = BlipForConditionalGeneration.from_pretrained(
    "Salesforce/blip-image-captioning-base")

def clean_and_resize(src_dir: str, dst_dir: str,
                     size: int = 1024) -> list:
    """统一分辨率、去重、剔除低质量图"""
    os.makedirs(dst_dir, exist_ok=True)
    kept = []
    seen_hashes = set()
    for name in os.listdir(src_dir):
        if not name.lower().endswith((".jpg", ".png")):
            continue
        path = os.path.join(src_dir, name)
        img = Image.open(path).convert("RGB")
        # 短边不足直接剔除，避免放大失真
        if min(img.size) < 512:
            continue
        h = hash(img.resize((8, 8)).tobytes())
        if h in seen_hashes:
            continue
        seen_hashes.add(h)
        img = img.resize((size, size), Image.LANCZOS)
        out = os.path.join(dst_dir, name)
        img.save(out, quality=95)
        kept.append(out)
    return kept

def auto_caption(image_path: str, trigger: str) -> str:
    """生成基础描述，人工再修正"""
    raw = Image.open(image_path).convert("RGB")
    inputs = processor(raw, return_tensors="pt")
    out = model.generate(**inputs, max_new_tokens=50)
    text = processor.decode(out[0], skip_special_tokens=True)
    return f"{trigger}, {text}"

from diffusers import StableDiffusionPipeline, DDPMScheduler
from peft import LoraConfig, get_peft_model
import torch

pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    torch_dtype=torch.float16).to("cuda")
unet = pipe.unet

lora_config = LoraConfig(
    r=16,                      # 秩：8-32 之间，人像常用 16
    lora_alpha=32,             # 缩放：通常为 r 的 2 倍
    target_modules=["to_q", "to_k", "to_v", "to_out.0"],
    lora_dropout=0.05,
    bias="none")
unet = get_peft_model(unet, lora_config)

optimizer = torch.optim.AdamW(
    unet.parameters(), lr=1e-4, weight_decay=1e-2)
scheduler = DDPMScheduler.from_pretrained(
    "runwayml/stable-diffusion-v1-5", subfolder="scheduler")

def train_step(batch, unet, vae, text_encoder, tokenizer,
               noise_scheduler, device="cuda"):
    images = batch["pixel_values"].to(device)
    captions = batch["captions"]

    latents = vae.encode(images).latent_dist.sample() * 0.18215
    noise = torch.randn_like(latents)
    timesteps = torch.randint(
        0, noise_scheduler.config.num_train_timesteps,
        (latents.shape[0],), device=device).long()
    noisy = noise_scheduler.add_noise(latents, noise, timesteps)

    tokens = tokenizer(
        captions, padding="max_length",
        max_length=tokenizer.model_max_length,
        truncation=True, return_tensors="pt").to(device)
    encoder_hidden = text_encoder(tokens.input_ids)[0]

    pred = unet(noisy, timesteps, encoder_hidden).sample
    loss = torch.nn.functional.mse_loss(pred, noise)
    loss.backward()
    return loss.item()

def evaluate(pipe, lora_path: str, prompts: list,
             seed: int = 42) -> dict:
    """固定 seed 和提示词，出图用于横向对比"""
    pipe.unload_lora_weights()
    pipe.load_lora_weights(lora_path)
    results = {}
    generator = torch.Generator("cuda").manual_seed(seed)
    for p in prompts:
        img = pipe(p, num_inference_steps=30,
                   guidance_scale=7.5,
                   generator=generator).images[0]
        results[p] = img
    return results

eval_prompts = [
    "ohwx person, close-up portrait, soft window light",
    "ohwx person, full body, outdoor, golden hour",
    "ohwx person, side profile, studio lighting, neutral background",
]