train_latent_text_to_image.py 8.88 KB
Newer Older
anton-l's avatar
anton-l committed
1
2
3
4
5
6
import argparse
import os

import torch
import torch.nn.functional as F

anton-l's avatar
anton-l committed
7
import bitsandbytes as bnb
anton-l's avatar
anton-l committed
8
9
10
import PIL.Image
from accelerate import Accelerator
from datasets import load_dataset
anton-l's avatar
anton-l committed
11
from diffusers import DDPMScheduler, LatentDiffusion, UNetLDMModel
anton-l's avatar
anton-l committed
12
13
14
15
16
17
18
from diffusers.hub_utils import init_git_repo, push_to_hub
from diffusers.optimization import get_scheduler
from diffusers.utils import logging
from torchvision.transforms import (
    CenterCrop,
    Compose,
    InterpolationMode,
anton-l's avatar
anton-l committed
19
    Normalize,
anton-l's avatar
anton-l committed
20
21
22
23
24
25
26
27
28
29
30
31
32
    RandomHorizontalFlip,
    Resize,
    ToTensor,
)
from tqdm.auto import tqdm


logger = logging.get_logger(__name__)


def main(args):
    accelerator = Accelerator(mixed_precision=args.mixed_precision)

anton-l's avatar
anton-l committed
33
34
    pipeline = LatentDiffusion.from_pretrained("fusing/latent-diffusion-text2im-large")
    pipeline.unet = None  # this model will be trained from scratch now
anton-l's avatar
anton-l committed
35
36
37
38
39
40
41
    model = UNetLDMModel(
        attention_resolutions=[4, 2, 1],
        channel_mult=[1, 2, 4, 4],
        context_dim=1280,
        conv_resample=True,
        dims=2,
        dropout=0,
anton-l's avatar
anton-l committed
42
        image_size=8,
anton-l's avatar
anton-l committed
43
44
45
46
47
48
49
50
51
52
53
54
55
        in_channels=4,
        model_channels=320,
        num_heads=8,
        num_res_blocks=2,
        out_channels=4,
        resblock_updown=False,
        transformer_depth=1,
        use_new_attention_order=False,
        use_scale_shift_norm=False,
        use_spatial_transformer=True,
        legacy=False,
    )
    noise_scheduler = DDPMScheduler(timesteps=1000, tensor_format="pt")
anton-l's avatar
anton-l committed
56
    optimizer = bnb.optim.Adam8bit(model.parameters(), lr=args.lr)
anton-l's avatar
anton-l committed
57
58
59
60
61
62
63

    augmentations = Compose(
        [
            Resize(args.resolution, interpolation=InterpolationMode.BILINEAR),
            CenterCrop(args.resolution),
            RandomHorizontalFlip(),
            ToTensor(),
anton-l's avatar
anton-l committed
64
            Normalize([0.5], [0.5]),
anton-l's avatar
anton-l committed
65
66
67
68
        ]
    )
    dataset = load_dataset(args.dataset, split="train")

anton-l's avatar
anton-l committed
69
70
71
    text_encoder = pipeline.bert.eval()
    vqvae = pipeline.vqvae.eval()

anton-l's avatar
anton-l committed
72
73
    def transforms(examples):
        images = [augmentations(image.convert("RGB")) for image in examples["image"]]
anton-l's avatar
anton-l committed
74
75
76
77
78
79
        text_inputs = pipeline.tokenizer(examples["caption"], padding="max_length", max_length=77, return_tensors="pt")
        with torch.no_grad():
            text_embeddings = accelerator.unwrap_model(text_encoder)(text_inputs.input_ids.cpu()).last_hidden_state
            images = 1 / 0.18215 * torch.stack(images, dim=0)
            latents = accelerator.unwrap_model(vqvae).encode(images.cpu()).mode()
        return {"images": images, "text_embeddings": text_embeddings, "latents": latents}
anton-l's avatar
anton-l committed
80
81
82
83
84
85
86
87
88
89
90

    dataset.set_transform(transforms)
    train_dataloader = torch.utils.data.DataLoader(dataset, batch_size=args.batch_size, shuffle=True)

    lr_scheduler = get_scheduler(
        "linear",
        optimizer=optimizer,
        num_warmup_steps=args.warmup_steps,
        num_training_steps=(len(train_dataloader) * args.num_epochs) // args.gradient_accumulation_steps,
    )

anton-l's avatar
anton-l committed
91
92
    model, text_encoder, vqvae, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
        model, text_encoder, vqvae, optimizer, train_dataloader, lr_scheduler
anton-l's avatar
anton-l committed
93
    )
anton-l's avatar
anton-l committed
94
95
    text_encoder = text_encoder.cpu()
    vqvae = vqvae.cpu()
anton-l's avatar
anton-l committed
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112

    if args.push_to_hub:
        repo = init_git_repo(args, at_init=True)

    # Train!
    is_distributed = torch.distributed.is_available() and torch.distributed.is_initialized()
    world_size = torch.distributed.get_world_size() if is_distributed else 1
    total_train_batch_size = args.batch_size * args.gradient_accumulation_steps * world_size
    max_steps = len(train_dataloader) // args.gradient_accumulation_steps * args.num_epochs
    logger.info("***** Running training *****")
    logger.info(f"  Num examples = {len(train_dataloader.dataset)}")
    logger.info(f"  Num Epochs = {args.num_epochs}")
    logger.info(f"  Instantaneous batch size per device = {args.batch_size}")
    logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_train_batch_size}")
    logger.info(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
    logger.info(f"  Total optimization steps = {max_steps}")

anton-l's avatar
anton-l committed
113
    global_step = 0
anton-l's avatar
anton-l committed
114
115
116
117
118
    for epoch in range(args.num_epochs):
        model.train()
        with tqdm(total=len(train_dataloader), unit="ba") as pbar:
            pbar.set_description(f"Epoch {epoch}")
            for step, batch in enumerate(train_dataloader):
anton-l's avatar
anton-l committed
119
120
121
122
                clean_latents = batch["latents"]
                noise_samples = torch.randn(clean_latents.shape).to(clean_latents.device)
                bsz = clean_latents.shape[0]
                timesteps = torch.randint(0, noise_scheduler.timesteps, (bsz,), device=clean_latents.device).long()
anton-l's avatar
anton-l committed
123

anton-l's avatar
anton-l committed
124
                # add noise onto the clean latents according to the noise magnitude at each timestep
anton-l's avatar
anton-l committed
125
                # (this is the forward diffusion process)
anton-l's avatar
anton-l committed
126
                noisy_latents = noise_scheduler.training_step(clean_latents, noise_samples, timesteps)
anton-l's avatar
anton-l committed
127
128
129

                if step % args.gradient_accumulation_steps != 0:
                    with accelerator.no_sync(model):
anton-l's avatar
anton-l committed
130
                        output = model(noisy_latents, timesteps, context=batch["text_embeddings"])
anton-l's avatar
anton-l committed
131
132
133
134
                        # predict the noise residual
                        loss = F.mse_loss(output, noise_samples)
                        loss = loss / args.gradient_accumulation_steps
                        accelerator.backward(loss)
anton-l's avatar
anton-l committed
135
                        optimizer.step()
anton-l's avatar
anton-l committed
136
                else:
anton-l's avatar
anton-l committed
137
                    output = model(noisy_latents, timesteps, context=batch["text_embeddings"])
anton-l's avatar
anton-l committed
138
139
140
141
142
143
144
145
146
147
                    # predict the noise residual
                    loss = F.mse_loss(output, noise_samples)
                    loss = loss / args.gradient_accumulation_steps
                    accelerator.backward(loss)
                    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
                    optimizer.step()
                    lr_scheduler.step()
                    optimizer.zero_grad()
                pbar.update(1)
                pbar.set_postfix(loss=loss.detach().item(), lr=optimizer.param_groups[0]["lr"])
anton-l's avatar
anton-l committed
148
                global_step += 1
anton-l's avatar
anton-l committed
149

anton-l's avatar
anton-l committed
150
        accelerator.wait_for_everyone()
anton-l's avatar
anton-l committed
151
152

        # Generate a sample image for visual inspection
anton-l's avatar
anton-l committed
153
        if accelerator.is_main_process:
anton-l's avatar
anton-l committed
154
155
            model.eval()
            with torch.no_grad():
anton-l's avatar
anton-l committed
156
                pipeline.unet = accelerator.unwrap_model(model)
anton-l's avatar
anton-l committed
157
158
159

                generator = torch.manual_seed(0)
                # run pipeline in inference (sample random noise and denoise)
anton-l's avatar
anton-l committed
160
161
162
                image = pipeline(
                    ["a clip art of a corgi"], generator=generator, eta=0.3, guidance_scale=6.0, num_inference_steps=50
                )
anton-l's avatar
anton-l committed
163
164
165

            # process image to PIL
            image_processed = image.cpu().permute(0, 2, 3, 1)
anton-l's avatar
anton-l committed
166
            image_processed = image_processed * 255.0
anton-l's avatar
anton-l committed
167
168
169
170
171
172
173
174
175
176
177
178
179
            image_processed = image_processed.type(torch.uint8).numpy()
            image_pil = PIL.Image.fromarray(image_processed[0])

            # save image
            test_dir = os.path.join(args.output_dir, "test_samples")
            os.makedirs(test_dir, exist_ok=True)
            image_pil.save(f"{test_dir}/{epoch:04d}.png")

            # save the model
            if args.push_to_hub:
                push_to_hub(args, pipeline, repo, commit_message=f"Epoch {epoch}", blocking=False)
            else:
                pipeline.save_pretrained(args.output_dir)
anton-l's avatar
anton-l committed
180
        accelerator.wait_for_everyone()
anton-l's avatar
anton-l committed
181
182
183
184
185


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Simple example of a training script.")
    parser.add_argument("--local_rank", type=int, default=-1)
anton-l's avatar
anton-l committed
186
187
    parser.add_argument("--dataset", type=str, default="fusing/dog_captions")
    parser.add_argument("--output_dir", type=str, default="ldm-text2image")
anton-l's avatar
anton-l committed
188
    parser.add_argument("--overwrite_output_dir", action="store_true")
anton-l's avatar
anton-l committed
189
190
    parser.add_argument("--resolution", type=int, default=128)
    parser.add_argument("--batch_size", type=int, default=1)
anton-l's avatar
anton-l committed
191
    parser.add_argument("--num_epochs", type=int, default=100)
anton-l's avatar
anton-l committed
192
    parser.add_argument("--gradient_accumulation_steps", type=int, default=16)
anton-l's avatar
anton-l committed
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
    parser.add_argument("--lr", type=float, default=1e-4)
    parser.add_argument("--warmup_steps", type=int, default=500)
    parser.add_argument("--push_to_hub", action="store_true")
    parser.add_argument("--hub_token", type=str, default=None)
    parser.add_argument("--hub_model_id", type=str, default=None)
    parser.add_argument("--hub_private_repo", action="store_true")
    parser.add_argument(
        "--mixed_precision",
        type=str,
        default="no",
        choices=["no", "fp16", "bf16"],
        help=(
            "Whether to use mixed precision. Choose"
            "between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."
            "and an Nvidia Ampere GPU."
        ),
    )

    args = parser.parse_args()
    env_local_rank = int(os.environ.get("LOCAL_RANK", -1))
    if env_local_rank != -1 and env_local_rank != args.local_rank:
        args.local_rank = env_local_rank

    main(args)