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.
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
