How to generate white background images instead of colorful?

[emadyounan]

Hi, I'm trying to train stable-diffusion-v1-5/stable-diffusion-v1-5 with custom data with LoraConfig to generate a type of 2D asset.
So I have a dataset of transparent PNG images, and I added to my prompt text for every image a white background after the caption
And in my script, I added BG_COLOR = (255, 255, 255)
and that's my code

BASE_MODEL_ID = "stable-diffusion-v1-5/stable-diffusion-v1-5"
DATASET_DIR = r"Dataset/10_postapo_icon_style"
OUTPUT_DIR = r"output_lora_test"
OUTPUT_NAME = "postapo_icon_lora_test.safetensors"

RESOLUTION = 512  # لو حصل OOM قلّلها لـ 448 أو 384
BATCH_SIZE = 1  # على كارت أحدث ممكن 2-4 (بيقلّل عدد الخطوات في الـ epoch، فظبط الـ epochs/LR حسبه)
LEARNING_RATE = 1e-4
NUM_EPOCHS = 15
LORA_RANK = 16
SAVE_EVERY = 200  # checkpoint كل كام خطوة
# None = تلقائي (بيشتغل لو الـ VRAM أقل من 12 GB). True/False لو عايز تفرضه
GRADIENT_CHECKPOINTING = None

# نفس لون الخلفية اللي اتكتبت بيه الكابشنات (BG_COLOR في auto_caption.py)
BG_COLOR = (255, 255, 255)
FALLBACK_CAPTION = "postapo_style, 2d, white background"  # لو صورة ملهاش .txt
IMAGE_EXTS = (".png", ".jpg", ".jpeg", ".webp")
# =========================================================

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
USE_AMP = DEVICE == "cuda"

VRAM_GB = 0.0
if USE_AMP:
    props = torch.cuda.get_device_properties(0)
    VRAM_GB = props.total_memory / 1024**3
    # bfloat16 لكروت Ampere وما فوق (RTX 30/40/50)، وfp16 للأقدم.
    # بنعتمد على compute capability لأن is_bf16_supported() بترجع True غلط على كروت قديمة
    AMP_DTYPE = torch.bfloat16 if props.major >= 8 else torch.float16
    print(f"GPU: {props.name} | {VRAM_GB:.1f} GB | compute capability {props.major}.{props.minor}")
else:
    AMP_DTYPE = torch.float32

# الأجزاء المجمّدة بنفس الـ dtype، وأوزان الـ LoRA نفسها fp32
WEIGHT_DTYPE = AMP_DTYPE
USE_GRAD_SCALER = AMP_DTYPE == torch.float16  # bf16 مش محتاج scaler
if GRADIENT_CHECKPOINTING is None:
    GRADIENT_CHECKPOINTING = 0 < VRAM_GB < 12

print(
    f"--- Starting LoRA Training on {DEVICE} ({WEIGHT_DTYPE}), "
    f"gradient checkpointing={GRADIENT_CHECKPOINTING} ---"
)


def load_square(path, size: int) -> Image.Image:
    """يلصق الشفافية على BG_COLOR، وبعدين يعمل padding لمربع بدل ما يمط الصورة."""
    img = Image.open(path)
    if img.mode in ("RGBA", "LA") or (img.mode == "P" and "transparency" in img.info):
        img = img.convert("RGBA")
        base = Image.new("RGBA", img.size, BG_COLOR + (255,))
        base.alpha_composite(img)
        img = base.convert("RGB")
    else:
        img = img.convert("RGB")

    w, h = img.size
    side = max(w, h)
    canvas = Image.new("RGB", (side, side), BG_COLOR)
    canvas.paste(img, ((side - w) // 2, (side - h) // 2))
    return canvas.resize((size, size), Image.LANCZOS)


class TextImageDataset(Dataset):
    def __init__(self, folder_path, resolution=512):
        folder = Path(folder_path)
        self.image_paths = sorted(
            p for p in folder.iterdir() if p.suffix.lower() in IMAGE_EXTS
        )
        self.resolution = resolution
        self.transform = transforms.Compose(
            [transforms.ToTensor(), transforms.Normalize([0.5], [0.5])]
        )

    def __len__(self):
        return len(self.image_paths)

    def __getitem__(self, idx):
        img_path = self.image_paths[idx]
        pixel_values = self.transform(load_square(img_path, self.resolution))

        txt_path = img_path.with_suffix(".txt")
        if txt_path.exists():
            caption = txt_path.read_text(encoding="utf-8").strip()
        else:
            caption = FALLBACK_CAPTION

        return {"pixel_values": pixel_values, "caption": caption}


# 3. تحميل القطع الأساسية
tokenizer = CLIPTokenizer.from_pretrained(BASE_MODEL_ID, subfolder="tokenizer")
text_encoder = CLIPTextModel.from_pretrained(
    BASE_MODEL_ID, subfolder="text_encoder"
).to(DEVICE, dtype=WEIGHT_DTYPE)
vae = AutoencoderKL.from_pretrained(BASE_MODEL_ID, subfolder="vae").to(
    DEVICE, dtype=WEIGHT_DTYPE
)
unet = UNet2DConditionModel.from_pretrained(BASE_MODEL_ID, subfolder="unet").to(
    DEVICE, dtype=WEIGHT_DTYPE
)
noise_scheduler = DDPMScheduler.from_pretrained(BASE_MODEL_ID, subfolder="scheduler")

for m in (unet, vae, text_encoder):
    m.requires_grad_(False)
text_encoder.eval()
vae.eval()

# 4. LoRA (بالطريقة الرسمية في diffusers: add_adapter بدل get_peft_model)
lora_config = LoraConfig(
    r=LORA_RANK,
    lora_alpha=LORA_RANK,
    target_modules=["to_q", "to_k", "to_v", "to_out.0"],
    lora_dropout=0.0,
    bias="none",
)
unet.add_adapter(lora_config)

# أوزان الـ LoRA لازم تفضل fp32 عشان التدريب يبقى مستقر
lora_params = [p for p in unet.parameters() if p.requires_grad]
for p in lora_params:
    p.data = p.data.float()
print(f"Trainable LoRA params: {sum(p.numel() for p in lora_params):,}")

if GRADIENT_CHECKPOINTING:
    unet.enable_gradient_checkpointing()

# 5. البيانات والـ Optimizer
dataset = TextImageDataset(DATASET_DIR, resolution=RESOLUTION)
print(f"Found {len(dataset)} images in {DATASET_DIR}")
dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)
optimizer = torch.optim.AdamW(lora_params, lr=LEARNING_RATE)
scaler = torch.amp.GradScaler("cuda", enabled=USE_GRAD_SCALER)

os.makedirs(OUTPUT_DIR, exist_ok=True)


def save_lora(save_dir: str) -> None:
    os.makedirs(save_dir, exist_ok=True)
    state_dict = convert_state_dict_to_diffusers(get_peft_model_state_dict(unet))
    StableDiffusionPipeline.save_lora_weights(
        save_directory=save_dir,
        unet_lora_layers=state_dict,
        weight_name=OUTPUT_NAME,
        safe_serialization=True,
    )


# 6. حلقة التدريب
unet.train()
step = 0
recent_losses = []

for epoch in range(NUM_EPOCHS):
    for batch in dataloader:
        images = batch["pixel_values"].to(DEVICE, dtype=WEIGHT_DTYPE)

        with torch.no_grad():
            latents = vae.encode(images).latent_dist.sample() * vae.config.scaling_factor
            input_ids = tokenizer(
                batch["caption"],
                padding="max_length",
                max_length=tokenizer.model_max_length,
                truncation=True,
                return_tensors="pt",
            ).input_ids.to(DEVICE)
            encoder_hidden_states = text_encoder(input_ids)[0]

        noise = torch.randn_like(latents)
        timesteps = torch.randint(
            0,
            noise_scheduler.config.num_train_timesteps,
            (latents.shape[0],),
            device=DEVICE,
        ).long()
        noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)

        with torch.autocast("cuda", dtype=AMP_DTYPE, enabled=USE_AMP):
            noise_pred = unet(noisy_latents, timesteps, encoder_hidden_states).sample
        loss = F.mse_loss(noise_pred.float(), noise.float())

        optimizer.zero_grad(set_to_none=True)
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()

        step += 1
        recent_losses.append(loss.item())
        if step % 10 == 0:
            avg = sum(recent_losses) / len(recent_losses)
            print(f"epoch {epoch + 1}/{NUM_EPOCHS}  step {step}  loss {avg:.4f}")
            recent_losses.clear()

        if step % SAVE_EVERY == 0:
            ckpt_dir = os.path.join(OUTPUT_DIR, f"checkpoint-{step}")
            save_lora(ckpt_dir)
            print(f"Saved checkpoint at step {step} -> {ckpt_dir}")

# 7. الحفظ النهائي
print("\nSaving final LoRA weights...")
save_lora(OUTPUT_DIR)
print(f"LoRA saved to: {OUTPUT_DIR}/{OUTPUT_NAME}")



pipe = StableDiffusionPipeline.from_pretrained(
    "stable-diffusion-v1-5/stable-diffusion-v1-5",
    torch_dtype=torch.float16,
    safety_checker=None,
    requires_safety_checker=False,
).to("cuda")


pipe.load_lora_weights("output_lora_test", weight_name="postapo_icon_lora_test.safetensors")

image = pipe("postapo_style, 2d, a rusty gas mask, isolated on white", num_inference_steps=40).images[0]
image.save("test.png")

But I got this image

thanks and regards

[TongyiDai]

Your load_square() function does composite transparent inputs onto white, so BG_COLOR itself is not an inference-time background control. The result means the LoRA/text conditioning has not learned “white background” strongly enough to override SD 1.5’s background prior.

I would debug it in this order:

1. Save a batch after load_square() and after reversing normalization. Verify that every training tensor really has an RGB white background—not only that the source PNG is transparent.
2. Use the exact same fixed phrase in every caption and at inference, e.g. “postapo_style, 2d icon, isolated on a plain white background.” In your example, training uses “white background” while inference says “isolated on white.” Consistency matters with a small dataset.
3. Check composition. If the object fills almost the whole 512×512 canvas, the loss contains little white-background area. Add examples with substantial, clean white negative space and varied objects.
4. Compare several checkpoints with a fixed seed. Dataset size, learning rate, rank, and total steps interact; do not judge only the final checkpoint. Also compare different LoRA scales at inference.
5. Add a negative prompt such as “colorful background, gradient background, scenery, texture, shadows on background” as a useful inference aid—not as a substitute for consistent data.

A UNet attention LoRA is good at learning style/concepts but is not a pixel-level background guarantee. If production output must be exactly #FFFFFF, the reliable solution is segmentation/alpha extraction followed by compositing onto white (or an explicit conditioning/mask workflow). Use the LoRA for the asset style and enforce the final background deterministically in post-processing.

[emadyounan]

I tested load_square(), and it saved images with a white background, and I added this to my caption
as an example, btw, I used only 100 images for training. How many images should I use
postapo_style, post-apocalyptic, 2D game icon, pure white background, a brown bomber jacket with a fur collar and cuffs

And this is what I'm using to generate images

import os
import torch
from diffusers import StableDiffusionPipeline

# ======================= الإعدادات =======================
BASE_MODEL_ID = "stable-diffusion-v1-5/stable-diffusion-v1-5"
LORA_PATH = r"output_lora_test"
LORA_NAME = "postapo_icon_lora_test.safetensors"

# قوة تأثير الـ LoRA (جرب قيم مختلفة من 0.7 إلى 1.2)
LORA_SCALE = 1.0

OUTPUT_DIR = "test_results"
os.makedirs(OUTPUT_DIR, exist_ok=True)

# إعدادات التوليد
NUM_IMAGES_PER_PROMPT = 2  # عدد الصور لكل prompt
STEPS = 40
GUIDANCE = 7.5
SIZE = 512

# ======================= قائمة الاختبارات =======================
TEST_PROMPTS = [
    # === الملابس العلوية ===
    "postapo_style, post-apocalyptic, 2d game icon, pure white background a blue t-shirt",
    "postapo_style, post-apocalyptic, 2d game icon, pure white background a brown leather jacket with zippers",
    "postapo_style, post-apocalyptic, 2d game icon, pure white background a gray hoodie",
    "postapo_style, post-apocalyptic, 2d game icon, pure white background a green military vest",
    "postapo_style, post-apocalyptic, 2d game icon, pure white background a white tank top",
    
    # === الملابس السفلية ===
    "postapo_style, post-apocalyptic, 2d game icon, pure white background brown cargo pants with pockets",
    "postapo_style, post-apocalyptic, 2d game icon, pure white background blue jeans with tears",
    "postapo_style, post-apocalyptic, 2d game icon, pure white background black tactical shorts",
    
    # === الأحذية ===
    "postapo_style, post-apocalyptic, 2d game icon, pure white background brown leather boots",
    "postapo_style, post-apocalyptic, 2d game icon, pure white background black combat boots",
    "postapo_style, post-apocalyptic, 2d game icon, pure white background gray sneakers",
    
    # === القبعات والإكسسوارات ===
    "postapo_style, post-apocalyptic, 2d game icon, pure white background a brown cap with star",
    "postapo_style, post-apocalyptic, 2d game icon, pure white background a gray beanie hat",
    "postapo_style, post-apocalyptic, 2d game icon, pure white background a blue scarf",
    "postapo_style, post-apocalyptic, 2d game icon, pure white background leather gloves",
    
    # === الأسلحة ===
    "postapo_style, post-apocalyptic, 2d game icon, pure white background a large axe with wooden handle",
    "postapo_style, post-apocalyptic, 2d game icon, pure white background a rusty sword",
    "postapo_style, post-apocalyptic, 2d game icon, pure white background a baseball bat with nails",
    "postapo_style, post-apocalyptic, 2d game icon, pure white background a crossbow",
    
    # === الحقائب والمعدات ===
    "postapo_style, post-apocalyptic, 2d game icon, pure white background a brown backpack with straps",
    "postapo_style, post-apocalyptic, 2d game icon, pure white background a blue duffel bag",
    "postapo_style, post-apocalyptic, 2d game icon, pure white background a first aid kit",
    "postapo_style, post-apocalyptic, 2d game icon, pure white background a water canteen",
    
    # === عناصر إضافية ===
    "postapo_style, post-apocalyptic, 2d game icon, pure white background sunglasses with cracked lens",
    "postapo_style, post-apocalyptic, 2d game icon, pure white background a wristwatch with leather strap",
    "postapo_style, post-apocalyptic, 2d game icon, pure white background a belt with pouches",
]

# ======================= الكود الرئيسي =======================
def main():
    print("=" * 60)
    print("🎮 Post-Apocalyptic LoRA Testing Script")
    print("=" * 60)
    
    device = "cuda" if torch.cuda.is_available() else "cpu"
    dtype = torch.float16 if device == "cuda" else torch.float32
    
    print(f"\n📥 Loading Base Model: {BASE_MODEL_ID}")
    print(f"💾 Loading LoRA: {LORA_PATH}/{LORA_NAME}")
    print(f"🎯 LoRA Scale: {LORA_SCALE}")
    print(f" Device: {device} ({dtype})\n")
    
    # تحميل الموديل
    pipe = StableDiffusionPipeline.from_pretrained(
        BASE_MODEL_ID,
        torch_dtype=dtype,
        safety_checker=None,
        requires_safety_checker=False,
    ).to(device)
    
    # تحميل الـ LoRA
    pipe.load_lora_weights(LORA_PATH, weight_name=LORA_NAME)
    print("✅ Model and LoRA loaded successfully!\n")
    
    # Negative Prompt ثابت
    negative_prompt = (
        "blurry, low quality, worst quality, watermark, text, signature, "
        "cropped, texture, shadows, floor, wall, ground, drop shadow, "
        "complex background, realistic, 3d, photograph, noise, grainy, "
        "distorted, deformed, ugly, bad anatomy, disfigured, extra limbs"
    )
    
    # غير الـ Negative Prompt بس، وسيب باقي الإعدادات كما هي
    # negative_prompt = (
    #     "realistic, 3d, photo, blurry, low quality, watermark, text, "
    #     "complex background, orange outlines, blue background, frame, border, "
    #     "weapons, guns, swords, accessories, extra items, "
    #     "ugly, deformed, bad anatomy, disfigured"
    # )

    # التوليد
    total_prompts = len(TEST_PROMPTS)
    
    for i, prompt in enumerate(TEST_PROMPTS, 1):
        print(f"\n[{i}/{total_prompts}] Generating: {prompt[:60]}...")
        
        try:
            # توليد الصور
            images = pipe(
                prompt=prompt,
                negative_prompt=negative_prompt,
                num_inference_steps=STEPS,
                guidance_scale=GUIDANCE,
                height=SIZE,
                width=SIZE,
                num_images_per_prompt=NUM_IMAGES_PER_PROMPT,
                cross_attention_kwargs={"scale": LORA_SCALE}
            ).images
            
            # حفظ الصور
            # استخراج اسم العنصر من الـ prompt للحفظ
            item_name = prompt.split("background, ")[-1].replace(" ", "_").replace(",", "")[:30]
            
            for j, img in enumerate(images, 1):
                filename = f"{i:03d}_{item_name}_v{j}.png"
                filepath = os.path.join(OUTPUT_DIR, filename)
                img.save(filepath)
                print(f"  ✅ Saved: {filepath}")
                
        except Exception as e:
            print(f"  ❌ Error: {e}")
            continue
    
    print("\n" + "=" * 60)
    print("🎉 Testing Complete!")
    print(f"📁 All images saved in: {OUTPUT_DIR}/")
    print("=" * 60)
    
    # إحصائيات
    print(f"\n📊 Summary:")
    print(f"   Total Prompts: {total_prompts}")
    print(f"   Images per Prompt: {NUM_IMAGES_PER_PROMPT}")
    print(f"   Expected Total Images: {total_prompts * NUM_IMAGES_PER_PROMPT}")
    print(f"   LoRA Scale Used: {LORA_SCALE}")
    print(f"   Steps: {STEPS}, Guidance: {GUIDANCE}")

if __name__ == "__main__":
    main()

That's the result:

samples of my dataset:

Thanks sir for your reply

[tk1475]

Adding to what @TongyiDai said, from doing a lot of this in production:

On how many images: 100 is fine for a style LoRA, so the count isn't your problem. You're asking an attention LoRA to guarantee a pixel-perfect white background, and that's the one thing it's weak at. It learns the postapo icon style well, but white background stays a soft hint that it will drop whenever the composition wants to. More images won't change that.

If you need clean white every time, don't get it from the LoRA at all. Generate the asset, then remove the background with rembg or BiRefNet and paste it onto white. For flat-background game icons that's the normal pipeline, and it's the only way to get an exact result.

If you want to stay single-pass, the things that actually help, in order:

• Same caption everywhere. You trained on "white background" but your first test said "isolated on white". Pick one phrase and use it in both. Your later "pure white background" prompts already fixed this.
• More white space in the training images. If the object fills the frame, the model barely sees any background to learn from.
• Sweep the LoRA scale from 0.7 to 1.1 with a fixed seed, and compare a few checkpoints instead of only the last one.

Short version: LoRA for the style, background removal in post. Then you stop fighting the model.