Personalize image generation and editing from a single reference image through inference-time LoRA optimization. Iteratively update model parameters based on visual similarity scores without training encoders or fine-tuning on multiple images.
Install with Codex or Claude Copy this prompt, paste it into Codex, Claude, or another assistant, and let it review the skill page and install it for you.
A direct command skips the review prompt. Inspect the source before running it.
Personalize image generation and editing from a single reference image through inference-time LoRA optimization. Iteratively update model parameters based on visual similarity scores without training encoders or fine-tuning on multiple images.
Core Concept
SISO (Single Image Iterative Subject-driven Generation and Editing) solves personalization using only one reference image through training-free inference-time optimization. Instead of pre-training encoders or fine-tuning on datasets, it iteratively optimizes Low-Rank Adaptation (LoRA) parameters based on how similar generated images are to the reference. This elegant approach enables plug-and-play personalization with any diffusion model architecture (SDXL, Flux, Sana) without architectural changes.
Architecture Overview
The SISO framework operates through a simple but effective optimization loop:
LoRA Parameter Initialization: Randomly initialize low-rank factors that plug into a frozen diffusion model, adding minimal parameter overhead
Iterative Optimization Loop: Generate images, compute visual similarity losses, and backpropagate gradients to update LoRA parameters
Loss Functions for Identity Preservation: DINO and IR embedding distances preserve subject identity while filtering background interference
Staged Inference Strategy: Optimize with simple prompts and minimal steps, then reuse optimized parameters with complex prompts and full denoising for quality
Implementation
LoRA Module and Parameter Management
LoRA enables parameter-efficient updates by decomposing weight changes into low-rank factors. This allows optimization of a frozen diffusion model without modifying millions of parameters.
# Initialize with small random values for stability
self
0.01
self
0.01
def
forward
self, x: torch.Tensor
"""Apply original module output plus low-rank update."""
self
self
self
self
self
return
def
inject_lora_into_model
diffusion_model: nn.Module, rank: int = 8
Dict
str
"""Inject LoRA adapters into key diffusion model layers."""
# Target attention and feedforward layers
for
in
if
isinstance
and
any
in
for
in
"to_q"
"to_k"
"to_v"
"to_out"
return
Loss Functions for Identity and Content Preservation
The identity loss combines DINO (excellent for instance-level similarity) and IR (item-level similarity) embeddings. For editing, an additional masked MSE loss preserves background.
import torch
import torch.nn.functional as F
from torchvision import models
classSubjectIdentityLoss(nn.Module):
"""Combined loss for subject identity preservation during generation."""def__init__(self, use_dino=True, use_ir=True):
super().__init__()
self.use_dino = use_dino
self.use_ir = use_ir
if use_dino:
# DINO model for instance-level similarity (effective for animals)self.dino_model = models.vgg16(pretrained=True)
self.dino_model.eval()
for param inself.dino_model.parameters():
param.requires_grad = Falseif use_ir:
# IR (image retrieval) model for item-level similarityself.ir_model = models.resnet50(pretrained=True)
self.ir_model.eval()
for param inself.ir_model.parameters():
param.requires_grad = Falsedefforward(
self, generated_image: torch.Tensor,
reference_image: torch.Tensor
) -> torch.Tensor:
"""Compute identity preservation loss."""
loss = 0.0ifself.use_dino:
# DINO captures instance-level similarity
gen_dino = self.dino_model(generated_image)
ref_dino = self.dino_model(reference_image)
dino_loss = F.cosine_embedding_loss(
gen_dino, ref_dino,
torch.ones(generated_image.shape[0]).to(generated_image.device)
)
loss += dino_loss
ifself.use_ir:
# IR assesses item-level similarity
gen_ir = self.ir_model(generated_image)
ref_ir = self.ir_model(reference_image)
ir_loss = F.mse_loss(gen_ir, ref_ir)
loss += ir_loss
return loss
classGenerationLoss(nn.Module):
"""Combined loss for image generation with identity preservation."""def__init__(self, w_identity=1.0, w_ir=1.0, w_prompt=0.5):
super().__init__()
self.identity_loss = SubjectIdentityLoss(use_dino=True, use_ir=True)
self.w_identity = w_identity
self.w_ir = w_ir
self.w_prompt = w_prompt
defforward(
self, generated: torch.Tensor, reference: torch.Tensor,
clip_score: float = 0.5) -> torch.Tensor:
"""Compute weighted loss combining identity and prompt alignment."""
identity_loss = self.identity_loss(generated, reference)
# Inverse CLIP score loss (higher CLIP score means lower loss)
prompt_loss = (1 - clip_score) * self.w_prompt
total_loss = self.w_identity * identity_loss + prompt_loss
return total_loss
classEditingLoss(nn.Module):
"""Loss for subject editing with background preservation."""def__init__(self, w_identity=1.0, w_bg=10.0):
super().__init__()
self.identity_loss = SubjectIdentityLoss(use_dino=True, use_ir=True)
self.w_identity = w_identity
self.w_bg = w_bg
defforward(
self, edited_image: torch.Tensor,
reference_image: torch.Tensor,
background_mask: torch.Tensor
) -> torch.Tensor:
"""Preserve subject identity while protecting background."""
identity_loss = self.identity_loss(edited_image, reference_image)
# Masked MSE loss: penalize background changes
bg_difference = F.mse_loss(
edited_image * background_mask,
reference_image * background_mask,
reduction='mean'
)
total_loss = (self.w_identity * identity_loss +
self.w_bg * bg_difference)
return total_loss
Iterative Optimization Loop
The core SISO algorithm iteratively generates images, computes losses, and updates LoRA parameters. Gradient normalization by loss magnitude ensures stable optimization across different loss scales.
defiterative_optimization_generation(
diffusion_model: nn.Module,
lora_modules: Dict[str, LoRA],
reference_image: torch.Tensor,
text_prompt: str,
num_iterations: int = 10,
learning_rate: float = 3e-4,
num_denoising_steps: int = 1,
improvement_threshold: float = 0.03,
patience: int = 7) -> Dict[str, LoRA]:
"""Iteratively optimize LoRA parameters for subject-driven generation."""# Collect LoRA parameters for optimization
lora_params = list(
param for lora in lora_modules.values()
for param in [lora.lora_a, lora.lora_b]
)
optimizer = torch.optim.Adam(lora_params, lr=learning_rate)
loss_fn = GenerationLoss()
best_loss = float('inf')
patience_counter = 0for iteration inrange(num_iterations):
# Stage 1: Generate image with current LoRA parameters
generated_image = diffusion_model.generate(
text_prompt, num_steps=num_denoising_steps,
lora_modules=lora_modules
)
# Compute loss
loss = loss_fn(generated_image, reference_image)
# Gradient-based update with loss magnitude normalization
optimizer.zero_grad()
loss.backward()
# Normalize gradients by loss magnitude for stable optimization
loss_magnitude = loss.detach().item()
if loss_magnitude > 0:
for param in lora_params:
if param.grad isnotNone:
param.grad = param.grad / (loss_magnitude + 1e-8)
optimizer.step()
# Early stopping: check for improvementif loss.item() < best_loss:
improvement = (best_loss - loss.item()) / (best_loss + 1e-8)
if improvement < improvement_threshold:
patience_counter += 1else:
patience_counter = 0
best_loss = loss.item()
else:
patience_counter += 1print(f"Iteration {iteration}: Loss = {loss.item():.6f}")
if patience_counter >= patience:
print(f"Early stopping at iteration {iteration}")
breakreturn lora_modules
defiterative_optimization_editing(
diffusion_model: nn.Module,
lora_modules: Dict[str, LoRA],
reference_image: torch.Tensor,
subject_mask: torch.Tensor,
text_prompt: str,
num_iterations: int = 10,
learning_rate: float = 3e-4,
num_denoising_steps: int = 3) -> Dict[str, LoRA]:
"""Optimize LoRA for subject editing with background preservation."""
lora_params = list(
param for lora in lora_modules.values()
for param in [lora.lora_a, lora.lora_b]
)
optimizer = torch.optim.Adam(lora_params, lr=learning_rate)
loss_fn = EditingLoss(w_identity=1.0, w_bg=10.0)
# Background mask is inverse of subject mask
background_mask = 1 - subject_mask
for iteration inrange(num_iterations):
# Generate edited version
edited_image = diffusion_model.generate(
text_prompt, num_steps=num_denoising_steps,
lora_modules=lora_modules
)
# Compute loss with background preservation
loss = loss_fn(edited_image, reference_image, background_mask)
optimizer.zero_grad()
loss.backward()
# Gradient normalization
loss_magnitude = loss.detach().item()
if loss_magnitude > 0:
for param in lora_params:
if param.grad isnotNone:
param.grad = param.grad / (loss_magnitude + 1e-8)
optimizer.step()
print(f"Iteration {iteration}: Edit Loss = {loss.item():.6f}")
return lora_modules
Two-Stage Inference Strategy
Optimize with simple prompts and minimal denoising steps for efficiency, then reuse the optimized LoRA parameters with complex prompts and full denoising for final quality.
defsiso_generation_pipeline(
diffusion_model: nn.Module,
reference_image: torch.Tensor,
simple_prompt: str,
final_prompt: str,
num_iterations_stage1: int = 10) -> torch.Tensor:
"""Two-stage SISO: optimize first, then generate with full quality."""# Inject LoRA into model
lora_modules = inject_lora_into_model(diffusion_model, rank=8)
# Stage 1: Optimize with simple prompt and minimal steps (efficiency)print("Stage 1: Optimization with simple prompt and 1 denoising step")
lora_modules = iterative_optimization_generation(
diffusion_model, lora_modules, reference_image,
text_prompt=simple_prompt,
num_iterations=num_iterations_stage1,
num_denoising_steps=1, # Minimal steps for speed
learning_rate=3e-4
)
# Stage 2: Generate final image with optimized LoRA (high quality)print("Stage 2: Final generation with complex prompt and full denoising")
final_image = diffusion_model.generate(
final_prompt,
num_steps=50, # Full denoising steps
lora_modules=lora_modules,
guidance_scale=7.5
)
return final_image
defsiso_editing_pipeline(
diffusion_model: nn.Module,
reference_image: torch.Tensor,
edit_prompt: str,
num_iterations: int = 10) -> torch.Tensor:
"""SISO editing: personalized subject with background preservation."""# Extract subject mask using Grounding DINO + SAM
subject_mask = extract_subject_mask(reference_image)
# Invert reference image to latent space using ReNoise
reference_latent = invert_to_latent_via_renoise(
diffusion_model, reference_image
)
# Inject and optimize LoRA
lora_modules = inject_lora_into_model(diffusion_model, rank=8)
lora_modules = iterative_optimization_editing(
diffusion_model, lora_modules,
reference_image, subject_mask,
text_prompt=edit_prompt,
num_iterations=num_iterations
)
# Generate edited image with optimized LoRA
edited_image = diffusion_model.generate(
edit_prompt,
num_steps=20,
lora_modules=lora_modules,
initial_latent=reference_latent
)
return edited_image
defextract_subject_mask(image: torch.Tensor) -> torch.Tensor:
"""Extract subject mask using Grounding DINO + SAM."""# Placeholder for Grounding DINO + SAM integration# In practice, use: https://github.com/IDEA-Research/Grounded-Segment-Anything
subject_mask = torch.zeros_like(image[:, :1, :, :])
return subject_mask
definvert_to_latent_via_renoise(
diffusion_model: nn.Module,
image: torch.Tensor
) -> torch.Tensor:
"""Invert image to latent space using ReNoise for editing."""# ReNoise inversion: iterative denoising-then-encoding# Placeholder for actual ReNoise implementation
latent = diffusion_model.encode(image)
return latent
Practical Guidance
When to use SISO:
You have a single reference image and want personalized generation/editing
You need flexible, fine-grained control over subject appearance and composition
Your diffusion model doesn't have built-in personalization (LoRA is architecture-agnostic)
You want interpretable, iterative results visible at each optimization step
You need to work across different model architectures (SDXL, Flux, Sana) without retraining
When NOT to use:
You have 10+ reference images (standard fine-tuning/DreamBooth more efficient)
You need immediate single-pass generation (optimization adds 10-30 seconds per image)
Your reference image is extremely noisy or low-quality (loss functions struggle)
You need to generate many images of the same subject (optimize once, reuse for efficiency)
Hyperparameter tuning:
Learning rate: 3e-4 is default; reduce to 1e-4 if optimization diverges, increase to 5e-4 if convergence is slow
LoRA rank: Default 8 balances capacity and efficiency; increase to 16 for complex subjects, reduce to 4 for simpler objects
Denoising steps (Stage 1): Single step (1) for efficiency; increase to 3-5 if subject is difficult to capture
Denoising steps (Stage 2): 50 steps for SDXL (good quality/speed balance), up to 70 for Flux, 20-30 for distilled models like SDXL-Turbo
Background weight (editing): Default 10.0; increase to 20 if background is critical, reduce to 5 if strict preservation is less important
Improvement threshold: 3% threshold triggers early stopping; reduce to 1% for faster convergence, increase to 5% for longer optimization
Common pitfalls:
Gradient explosion with high loss values: Always normalize gradients by loss magnitude; skipping this causes training instability
Mode collapse: Ensemble DINO + IR losses prevent collapse; using only one loss function increases risk
Over-optimization in Stage 1: Optimize too long with simple prompts and model memorizes artificial features; 10 iterations is usually sufficient
Background preservation failure in editing: Insufficient background weight or poor subject mask extraction; validate mask quality before editing
Diffusion step mismatch: Optimizing with 1 step but inferring with 50 steps can cause distribution shift; use consistent denoising schedules or backprop through multiple steps during optimization
Reference
Base architecture: SDXL, Flux Schnell, Sana diffusion models; LoRA from Microsoft's LoRA adaptation
Loss components: DINO embeddings for instance similarity (DINOv2 vision transformers), ResNet-50 features for item-level IR, CLIP-T for prompt alignment
Evaluation metrics: FID (image quality), DINO score (identity preservation), IR score (object consistency), LPIPS (background preservation in editing), CLIP-T (prompt adherence), user studies (naturalness and alignment)
Optimization technique: Gradient-normalized Adam with loss magnitude stabilization
Related work: DreamBooth (per-subject fine-tuning), TextualInversion (token learning), LoRA fine-tuning, ReNoise inversion, Grounding DINO + SAM for segmentation