Single Image Iterative Subject Generation

---
name: single-image-iterative-subject-generation
title: "Single Image Iterative Subject-driven Generation and Editing"
version: 0.0.2
engine: skillxiv-v0.0.2-claude-opus-4.6
license: MIT
url: "https://arxiv.org/abs/2503.16025"
keywords: [Image Generation, Personalization, LoRA Optimization, Diffusion Models, Subject-Driven Synthesis]
description: "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.

```python
import torch
import torch.nn as nn
from typing import Dict, Tuple

class LoRA(nn.Module):
    """Low-rank adaptation module for parameter-efficient fine-tuning."""

    def __init__(self, original_module: nn.Linear, rank: int = 8, alpha: float = 1.0):
        super().__init__()
        self.original_module = original_module
        self.rank = rank
        self.alpha = alpha

        # Low-rank factors
        self.lora_a = nn.Parameter(
            torch.zeros(original_module.in_features, rank)
        )
        self.lora_b = nn.Parameter(
            torch.zeros(rank, original_module.out_features)
        )

        # Initialize with small random values for stability
        nn.init.normal_(self.lora_a, std=0.01)
        nn.init.normal_(self.lora_b, std=0.01)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Apply original module output plus low-rank update."""
        original_out = self.original_module(x)
        lora_update = (self.alpha / self.rank) * (x @ self.lora_a @ self.lora_b)
        return original_out + lora_update

def inject_lora_into_model(
    diffusion_model: nn.Module, rank: int = 8
) -> Dict[str, LoRA]:
    """Inject LoRA adapters into key diffusion model layers."""
    lora_modules = {}

    # Target attention and feedforward layers
    for name, module in diffusion_model.named_modules():
        if isinstance(module, nn.Linear) and any(
            target in name for target in ["to_q", "to_k", "to_v", "to_out"]
        ):
            lora = LoRA(module, rank=rank)
            lora_modules[name] = lora

    return lora_modules
```

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

```python
import torch
import torch.nn.functional as F
from torchvision import models

class SubjectIdentityLoss(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 in self.dino_model.parameters():
                param.requires_grad = False

        if use_ir:
            # IR (image retrieval) model for item-level similarity
            self.ir_model = models.resnet50(pretrained=True)
            self.ir_model.eval()
            for param in self.ir_model.parameters():
                param.requires_grad = False

    def forward(
        self, generated_image: torch.Tensor,
        reference_image: torch.Tensor
    ) -> torch.Tensor:
        """Compute identity preservation loss."""
        loss = 0.0

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

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

class GenerationLoss(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

    def forward(
        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

class EditingLoss(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

    def forward(
        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.

```python
def iterative_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 = 0

    for iteration in range(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 is not None:
                    param.grad = param.grad / (loss_magnitude + 1e-8)

        optimizer.step()

        # Early stopping: check for improvement
        if loss.item() < best_loss:
            improvement = (best_loss - loss.item()) / (best_loss + 1e-8)
            if improvement < improvement_threshold:
                patience_counter += 1
            else:
                patience_counter = 0
            best_loss = loss.item()
        else:
            patience_counter += 1

        print(f"Iteration {iteration}: Loss = {loss.item():.6f}")

        if patience_counter >= patience:
            print(f"Early stopping at iteration {iteration}")
            break

    return lora_modules

def iterative_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 in range(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 is not None:
                    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.

```python
def siso_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

def siso_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

def extract_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

def invert_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