image-gen/run.py

# Copyright 2026-present the HuggingFace Inc. team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Main entry point for image generation method comparison experiments.

Based on https://github.com/huggingface/diffusers/blob/bbbcdd87bd9d960fa372663a50b9edbdcb1391c6/examples/dreambooth/train_dreambooth_lora_flux2_klein.py
"""

import argparse
import copy
import datetime as dt
import json
import os
import sys
import time
from collections.abc import Callable
from contextlib import AbstractContextManager, nullcontext
from functools import partial
from typing import Any, Optional

import huggingface_hub
import torch
from diffusers.training_utils import (
    compute_density_for_timestep_sampling,
    compute_loss_weighting_for_sd3,
    offload_models,
)
from torch.amp import GradScaler, autocast
from tqdm import tqdm
from transformers import set_seed
from utils import (
    FILE_NAME_TRAIN_PARAMS,
    TrainConfig,
    TrainResult,
    TrainStatus,
    get_artifact_stem,
    get_base_model_info,
    get_dataset_info,
    get_dino_embeddings,
    get_dino_encoder,
    get_file_size,
    get_optimizer_and_scheduler,
    get_peft_branch,
    get_pipeline,
    get_sample_image_save_dir,
    get_torch_dtype,
    get_train_config,
    init_accelerator,
    log_results,
    upload_checkpoint_to_bucket,
    upload_images_to_bucket,
    validate_experiment_path,
)

from data import get_train_valid_test_datasets
from peft import PeftConfig, PeftModel
from peft.utils import CONFIG_NAME, infer_device


os.environ["TORCHINDUCTOR_FORCE_DISABLE_CACHES"] = "1"


def get_sigmas(timesteps, noise_scheduler, *, step_indices, n_dim, dtype):
    device = "cpu"
    sigmas = noise_scheduler.sigmas.to(device=device, dtype=dtype)
    schedule_timesteps = noise_scheduler.timesteps.to(device)
    timesteps = timesteps.to(device)

    sigma = sigmas[step_indices].flatten()
    while len(sigma.shape) < n_dim:
        sigma = sigma.unsqueeze(-1)
    return sigma


class DummyGradScaler:
    def scale(self, loss):
        return loss

    def unscale_(self, optimizer):
        pass

    def step(self, optimizer):
        optimizer.step()

    def update(self):
        pass


def precompute_prompt_caches(
    pipeline, train_prompts: list[str], eval_prompts: list[str], device_type: str, train_config: TrainConfig
) -> tuple[torch.Tensor, torch.Tensor, dict[str, torch.Tensor]]:
    """Precompute all prompt embeds needed during the run in a single pass over the text encoder.

    This is mainly to save memory so that we don't need to either unload the text encoder during evaluation, which
    requires extra VRAM, or compute the embeddings on CPU, which is slow. Once this is done, we can completely remove
    the text encoder.

    The training prompts are returned as concatenated tensors on the accelerator. The eval prompts are returned as a
    dict mapping each prompt to its embedding on the CPU. The eval cache always includes the empty prompt: when the
    pipeline applies classifier-free guidance (guidance_scale > 1 on a non-distilled checkpoint), it requires embeddings
    for the empty negative prompt, which it would otherwise compute with the text encoder at generation time.

    """
    prompt_embeds_cache = []
    text_ids_cache = []
    eval_prompt_cache: dict[str, torch.Tensor] = {}
    with torch.no_grad(), offload_models(pipeline.text_encoder, device=device_type, offload=True):
        for prompt in train_prompts:
            prompt_embeds, text_ids = pipeline.encode_prompt(
                prompt=prompt,
                max_sequence_length=train_config.max_sequence_length,
                text_encoder_out_layers=train_config.text_encoder_out_layers,
            )
            prompt_embeds_cache.append(prompt_embeds)
            text_ids_cache.append(text_ids)
        empty_prompt = [""]  # e.g. for negative prompt
        for prompt in eval_prompts + empty_prompt:
            prompt_embeds, _ = pipeline.encode_prompt(
                prompt=prompt,
                max_sequence_length=train_config.max_sequence_length,
                text_encoder_out_layers=train_config.text_encoder_out_layers,
            )
            eval_prompt_cache[prompt] = prompt_embeds.to("cpu")
    return (
        (torch.cat(prompt_embeds_cache, dim=0) if prompt_embeds_cache else torch.tensor([])).to(device_type),
        (torch.cat(text_ids_cache, dim=0) if text_ids_cache else torch.tensor([])).to(device_type),
        eval_prompt_cache,
    )


def precompute_latent_cache(
    *,
    pipeline,
    vae,
    pixel_values: list[torch.Tensor],
    train_config: TrainConfig,
    device_type: str,
) -> torch.Tensor:
    latents_cache = []
    latents_bn_mean = vae.bn.running_mean.view(1, -1, 1, 1)
    latents_bn_std = torch.sqrt(vae.bn.running_var.view(1, -1, 1, 1) + vae.config.batch_norm_eps)
    with torch.no_grad(), offload_models(vae, device=device_type, offload=True):
        latents_bn_mean = latents_bn_mean.to(vae.device)
        latents_bn_std = latents_bn_std.to(vae.device)
        for i in range(0, len(pixel_values), train_config.batch_size):
            pixel_values_batch = torch.stack(pixel_values[i : i + train_config.batch_size]).to(
                device=vae.device, dtype=get_torch_dtype(train_config.dtype)
            )
            latents = vae.encode(pixel_values_batch).latent_dist.mode()
            latents = pipeline._patchify_latents(latents)
            latents = (latents - latents_bn_mean) / latents_bn_std
            latents_cache.append(latents.to(device_type))
    return torch.cat(latents_cache, dim=0)


def _generate_images(
    pipeline, *, generator, prompts: list[str], prompt_cache: dict[str, torch.Tensor], config: TrainConfig
):
    device = pipeline.transformer.device
    prompt_embeds = torch.cat([prompt_cache[prompt] for prompt in prompts], dim=0).to(device)
    negative_prompt_embeds = prompt_cache[""].to(device).expand(len(prompts), -1, -1)
    outputs = pipeline(
        prompt_embeds=prompt_embeds,
        negative_prompt_embeds=negative_prompt_embeds,
        num_inference_steps=config.num_inference_steps,
        guidance_scale=config.guidance_scale,
        height=config.resolution,  # hard-code square
        width=config.resolution,
        generator=generator,
        output_type="pil",
    )
    return outputs


@torch.inference_mode()
def evaluate(
    *,
    pipeline,
    ds_eval,
    processor,
    dino_model,
    prompt_cache: dict[str, torch.Tensor],
    config: TrainConfig,
    num_repeats: int = 1,
) -> float:
    with offload_models(pipeline.vae, device=pipeline.transformer.device, offload=True):
        # avoid reusing same seed as in training, which would bias samples toward memorized results
        seed = config.seed + 100_000
        generator = torch.Generator(device=pipeline.transformer.device).manual_seed(seed)
        cosine_sim_scores = []
        iter_ = range(num_repeats) if num_repeats <= 1 else tqdm(range(num_repeats))
        for _ in iter_:
            generated_images = []
            reference_images = []
            batch_size = config.batch_size_eval

            for i in range(0, len(ds_eval), batch_size):
                sliced = [ds_eval[j] for j in range(i, min(i + batch_size, len(ds_eval)))]
                prompts = [sample["prompt"] for sample in sliced]
                outputs = _generate_images(
                    pipeline, generator=generator, prompts=prompts, prompt_cache=prompt_cache, config=config
                )
                generated_images.extend(outputs.images)
                reference_images.extend([sample["raw_image"] for sample in sliced])
                if i + batch_size >= len(ds_eval):
                    break

            generated_embeddings = get_dino_embeddings(generated_images, processor, dino_model, batch_size=batch_size)
            reference_embeddings = get_dino_embeddings(reference_images, processor, dino_model, batch_size=batch_size)
            cosine_sim = (generated_embeddings * reference_embeddings).sum(dim=-1)
            cosine_sim_scores.append(cosine_sim.mean().item())
        mean_sim = sum(cosine_sim_scores) / num_repeats
    return mean_sim


@torch.inference_mode()
def measure_drift(
    *, pipeline, processor, dino_model, prompt_cache: dict[str, torch.Tensor], config: TrainConfig
) -> float:
    # Measure the drift as 1 - the cosine similarity of the images generated by the base model vs the model with the
    # trained adapter. The prompts are unrelated to the concept, so we expect the similarity to be high, hence the drift
    # to be low.
    if not isinstance(pipeline.transformer, PeftModel):
        # in case of full fine-tuning, the adapter cannot be disabled and thus the drift cannot be measured, return
        # dummy value
        return float("nan")

    batch_size = config.batch_size_eval
    prompts = config.drift_image_prompts
    pbar = tqdm(total=len(prompts) * 2)
    with offload_models(pipeline.vae, device=pipeline.transformer.device, offload=True):
        # without adapter
        # avoid reusing same seed as in training, which would bias samples toward memorized results
        seed = config.seed + 100_000_000
        generator = torch.Generator(device=pipeline.transformer.device).manual_seed(seed)
        generated_base = []
        with pipeline.transformer.disable_adapter():
            for i in range(0, len(prompts), batch_size):
                prompt_batch = prompts[i : i + batch_size]
                outputs = _generate_images(
                    pipeline, generator=generator, prompts=prompt_batch, prompt_cache=prompt_cache, config=config
                )
                generated_base.extend(outputs.images)
                pbar.update(1)

        # with adapter
        # avoid reusing same seed as in training, which would bias samples toward memorized results
        seed = config.seed + 100_000_000
        generator = torch.Generator(device=pipeline.transformer.device).manual_seed(seed)
        generated_adapter = []
        for i in range(0, len(prompts), batch_size):
            prompt_batch = prompts[i : i + batch_size]
            outputs = _generate_images(
                pipeline, generator=generator, prompts=prompt_batch, prompt_cache=prompt_cache, config=config
            )
            generated_adapter.extend(outputs.images)
            pbar.update(1)

    # calculate drift
    generated_embeddings = get_dino_embeddings(generated_adapter, processor, dino_model, batch_size=batch_size)
    reference_embeddings = get_dino_embeddings(generated_base, processor, dino_model, batch_size=batch_size)
    cosine_sim = (generated_embeddings * reference_embeddings).sum(dim=-1)  # dino embeddings are L2-normalized
    drift = (1 - cosine_sim.mean().item()) / 2.0  # cos sim is in [-1, 1], normalized to [0, 1]
    return drift


def train(
    *,
    pipeline,
    train_config: TrainConfig,
    accelerator_memory_init: int,
    is_adalora: bool,
    is_adamss: bool,
    print_verbose: Callable[..., None],
    device_type: str,
) -> tuple[TrainResult, dict[str, torch.Tensor]]:
    accelerator_memory_allocated_log = []
    accelerator_memory_reserved_log = []
    losses = []
    durations = []
    metrics = []
    total_samples = 0

    train_dataset, valid_dataset, test_dataset = get_train_valid_test_datasets(
        train_config=train_config, print_fn=print_verbose
    )
    train_size_base = len(train_dataset["prompts"])
    gen = torch.Generator(device=device_type).manual_seed(train_config.seed)
    train_indices = torch.cat(
        [torch.randperm(train_size_base, generator=gen, device=device_type) for _ in range(train_dataset["repeats"])]
    )
    if train_config.max_steps > len(train_indices):
        raise ValueError(
            f"max_steps is too high ({train_config.max_steps}), there are only {len(train_indices)} training samples"
        )

    processor, dino_model = get_dino_encoder(train_config.dino_model_id, train_config.dino_image_size)

    torch_accelerator_module = getattr(torch, device_type, torch.cuda)
    if train_config.use_amp:
        grad_scaler: GradScaler | DummyGradScaler = GradScaler(device=device_type)
        autocast_ctx: Callable[[], AbstractContextManager[Any]] = partial(autocast, device_type=device_type)
    else:
        grad_scaler = DummyGradScaler()
        autocast_ctx = nullcontext

    vae = pipeline.vae  # CPU
    transformer = pipeline.transformer
    noise_scheduler_copy = copy.deepcopy(pipeline.scheduler)  # prevent mutating it
    optimizer, lr_scheduler = get_optimizer_and_scheduler(
        transformer,
        optimizer_type=train_config.optimizer_type,
        max_steps=train_config.max_steps,
        lr_scheduler_arg=train_config.lr_scheduler,
        **train_config.optimizer_kwargs,
    )

    if hasattr(transformer, "get_nb_trainable_parameters"):
        num_trainable_params, num_params = transformer.get_nb_trainable_parameters()
    else:
        num_params = sum(param.numel() for param in transformer.parameters())
        num_trainable_params = sum(param.numel() for param in transformer.parameters() if param.requires_grad)
    print_verbose(
        f"trainable params: {num_trainable_params:,d} || all params: {num_params:,d} || "
        f"trainable: {100 * num_trainable_params / num_params:.4f}%"
    )

    status = TrainStatus.FAILED
    tic_train = time.perf_counter()
    eval_time = 0.0
    error_msg = ""

    # pre-compute, since they don't change during training and we can keep the text encoder and VAE offloaded; the
    # eval prompts are also encoded now, so that the text encoder is not needed anymore afterwards
    eval_prompts = (
        [sample["prompt"] for sample in valid_dataset]
        + [sample["prompt"] for sample in test_dataset]
        + list(train_config.drift_image_prompts)
        + list(train_config.sample_image_prompts)
    )
    prompt_embeds_cache, text_ids_cache, prompt_cache = precompute_prompt_caches(
        pipeline,
        train_prompts=train_dataset["prompts"],
        eval_prompts=eval_prompts,
        device_type=device_type,
        train_config=train_config,
    )
    latents_cache = precompute_latent_cache(
        pipeline=pipeline,
        vae=vae,
        pixel_values=train_dataset["pixel_values"],
        train_config=train_config,
        device_type=device_type,
    )

    # All prompts used in this run are cached now, so the text encoder is no longer needed. Drop it from the pipeline to
    # free memory
    pipeline.text_encoder = None
    transformer.to(device_type)

    torch_accelerator_module.empty_cache()
    torch_accelerator_module.reset_peak_memory_stats()
    accelerator_memory_max_train = 0
    try:
        torch_accelerator_module.reset_peak_memory_stats()
        pbar = tqdm(range(1, train_config.max_steps + 1))
        for step in pbar:
            tic = time.perf_counter()
            i_start = (step - 1) * train_config.batch_size
            i_stop = min(step * train_config.batch_size, len(train_indices))
            batch_indices = train_indices[i_start:i_stop].to(device=latents_cache.device, dtype=torch.long)
            latents = latents_cache.index_select(0, batch_indices)
            prompt_embeds = prompt_embeds_cache.index_select(0, batch_indices)
            text_ids = text_ids_cache.index_select(0, batch_indices)

            current_batch_size = latents.shape[0]
            total_samples += current_batch_size

            model_input_ids = pipeline._prepare_latent_ids(latents).to(latents.device)
            noise = torch.randn_like(latents, generator=gen)

            u = compute_density_for_timestep_sampling(
                weighting_scheme=train_config.weighting_scheme,
                batch_size=current_batch_size,
                logit_mean=train_config.logit_mean,
                logit_std=train_config.logit_std,
                mode_scale=train_config.mode_scale,
            )
            indices = (u * noise_scheduler_copy.config.num_train_timesteps).long()
            timesteps = noise_scheduler_copy.timesteps[indices].to(device=latents.device)
            # Add noise according to flow matching. zt = (1 - texp) * x + texp * z1
            sigmas = get_sigmas(
                timesteps, noise_scheduler_copy, step_indices=indices, n_dim=latents.ndim, dtype=latents.dtype
            ).to(device_type)
            noisy_latents = (1.0 - sigmas) * latents + sigmas * noise
            # [B, C, H, W] -> [B, H*W, C]
            packed_noisy_latents = pipeline._pack_latents(noisy_latents)

            # handle guidance
            if transformer.config.guidance_embeds:
                guidance = torch.full([1], train_config.guidance_scale, device=device_type)
                guidance = guidance.expand(current_batch_size)
            else:
                guidance = None

            optimizer.zero_grad(set_to_none=True)
            with autocast_ctx():
                model_pred = transformer(
                    hidden_states=packed_noisy_latents,
                    timestep=timesteps / 1000,
                    guidance=guidance,
                    encoder_hidden_states=prompt_embeds,
                    txt_ids=text_ids,  # B, text_seq_len, 4
                    img_ids=model_input_ids,  # B, image_seq_len, 4
                    return_dict=False,
                )[0]
                model_pred = model_pred[:, : packed_noisy_latents.size(1)]
                model_pred = pipeline._unpack_latents_with_ids(model_pred, model_input_ids)
                # these weighting schemes use a uniform timestep sampling and instead post-weight the loss
                weighting = compute_loss_weighting_for_sd3(train_config.weighting_scheme, sigmas=sigmas)
                target = noise - latents
                loss = torch.mean(
                    (weighting.float() * (model_pred.float() - target.float()) ** 2).reshape(target.shape[0], -1), 1
                )
                loss = loss.mean()

            grad_scaler.scale(loss).backward()
            if train_config.grad_norm_clip:
                grad_scaler.unscale_(optimizer)
                torch.nn.utils.clip_grad_norm_(transformer.parameters(), train_config.grad_norm_clip)
            grad_scaler.step(optimizer)
            grad_scaler.update()
            lr_scheduler.step()

            if is_adalora or is_adamss:
                transformer.base_model.update_and_allocate(step)

            loss = loss.detach()
            losses.append(loss)
            pbar.set_postfix({"loss": loss.item()})

            accelerator_memory_allocated_log.append(
                torch_accelerator_module.memory_allocated() - accelerator_memory_init
            )
            accelerator_memory_reserved_log.append(
                torch_accelerator_module.memory_reserved() - accelerator_memory_init
            )
            toc = time.perf_counter()
            durations.append(toc - tic)

            if step % train_config.eval_steps == 0:
                # Measure max memory _before_ executing the eval loop and reset stats _after_ the eval loop. This way
                # the extra memory required for evaluation is not included in the max memory statistic. We want to
                # measure only the training memory, as the eval requires extra memory (DINO model) not caused by the
                # PEFT method.
                accelerator_memory_max_train = max(
                    accelerator_memory_max_train,
                    torch_accelerator_module.max_memory_reserved() - accelerator_memory_init,
                )

                tic_eval = time.perf_counter()
                loss_avg = sum(losses[-train_config.eval_steps :]) / train_config.eval_steps
                loss_avg = loss_avg.item()
                memory_allocated_avg = (
                    sum(accelerator_memory_allocated_log[-train_config.eval_steps :]) / train_config.eval_steps
                )
                memory_reserved_avg = (
                    sum(accelerator_memory_reserved_log[-train_config.eval_steps :]) / train_config.eval_steps
                )
                dur_train = sum(durations[-train_config.eval_steps :])

                transformer.eval()
                valid_similarity = evaluate(
                    pipeline=pipeline,
                    ds_eval=valid_dataset,
                    processor=processor,
                    dino_model=dino_model,
                    prompt_cache=prompt_cache,
                    config=train_config,
                )
                transformer.train()

                toc_eval = time.perf_counter()
                dur_eval = toc_eval - tic_eval
                eval_time += dur_eval
                elapsed = time.perf_counter() - tic_train

                metrics.append(
                    {
                        "step": step,
                        "valid dino_similarity": valid_similarity,
                        "train loss": loss_avg,
                        "train samples": total_samples,
                        "train time": dur_train,
                        "eval time": dur_eval,
                        "mem allocated avg": memory_allocated_avg,
                        "mem reserved avg": memory_reserved_avg,
                        "elapsed time": elapsed,
                    }
                )

                log_dict = {
                    "step": f"{step:4d}",
                    "samples": f"{total_samples:5d}",
                    "lr": f"{lr_scheduler.get_last_lr()[0]:.2e}",
                    "loss avg": f"{loss_avg:.4f}",
                    "valid sim": f"{valid_similarity:.4f}",
                    "train time": f"{dur_train:.1f}s",
                    "eval time": f"{dur_eval:.1f}s",
                    "mem allocated": f"{memory_allocated_avg:.0f}",
                    "mem reserved": f"{memory_reserved_avg:.0f}",
                    "elapsed time": f"{elapsed // 60:.0f}min {elapsed % 60:.0f}s",
                }
                print_verbose(json.dumps(log_dict))

                torch_accelerator_module.empty_cache()
                torch_accelerator_module.reset_peak_memory_stats()

        accelerator_memory_max_train = max(
            accelerator_memory_max_train,
            torch_accelerator_module.max_memory_reserved() - accelerator_memory_init,
        )
        print_verbose(f"Training finished after {train_config.max_steps} steps, evaluation on test set follows.")
        transformer.eval()
        test_similarity = evaluate(
            pipeline=pipeline,
            ds_eval=test_dataset,
            processor=processor,
            dino_model=dino_model,
            prompt_cache=prompt_cache,
            config=train_config,
            num_repeats=3,
        )
        print_verbose("Calculating drift.")
        test_drift = measure_drift(
            pipeline=pipeline,
            processor=processor,
            dino_model=dino_model,
            prompt_cache=prompt_cache,
            config=train_config,
        )
        metrics.append(
            {
                "step": step,
                "test dino_similarity": test_similarity,
                "drift": test_drift,
                "train loss": (sum(losses[-train_config.eval_steps :]) / train_config.eval_steps).item(),
                "train samples": total_samples,
            }
        )
        print_verbose(f"Test DINOv2 similarity: {test_similarity:.4f}")
        print_verbose(f"Test drift:             {test_drift:.4f}")

    except KeyboardInterrupt:
        print_verbose("canceled training")
        status = TrainStatus.CANCELED
        error_msg = "manually canceled"
    except torch.OutOfMemoryError as exc:
        print_verbose("out of memory error encountered")
        status = TrainStatus.CANCELED
        error_msg = str(exc)
    except Exception as exc:
        print_verbose(f"encountered an error: {exc}")
        status = TrainStatus.CANCELED
        error_msg = str(exc)

    toc_train = time.perf_counter()
    train_time = toc_train - tic_train - eval_time

    if status != TrainStatus.CANCELED:
        status = TrainStatus.SUCCESS
    train_result = TrainResult(
        status=status,
        train_time=train_time,
        accelerator_memory_reserved_log=accelerator_memory_reserved_log,
        accelerator_memory_max_train=accelerator_memory_max_train,
        losses=[loss.item() for loss in losses],
        metrics=metrics,
        error_msg=error_msg,
        num_trainable_params=num_trainable_params,
        num_total_params=num_params,
    )
    # the prompt cache is returned so that generate_sample_images can reuse it, the text encoder is gone by now
    return train_result, prompt_cache


@torch.inference_mode()
def generate_sample_images(
    *,
    pipeline,
    train_config,
    prompt_cache: dict[str, torch.Tensor],
    sample_image_dir: str,
    file_stem: str,
) -> None:
    target_device = pipeline.transformer.device
    with offload_models(pipeline.vae, device=target_device, offload=True):
        # avoid reusing same seed as in training, which would bias samples toward memorized results
        seed = train_config.seed + 100_000
        generator = torch.Generator(device=target_device).manual_seed(seed)
        pbar = tqdm(
            enumerate(train_config.sample_image_prompts, start=1), total=len(train_config.sample_image_prompts)
        )
        for idx, prompt in pbar:
            image_path = os.path.join(sample_image_dir, f"{file_stem}_{idx:02d}.png")
            outputs = _generate_images(
                pipeline, generator=generator, prompts=[prompt], prompt_cache=prompt_cache, config=train_config
            )
            outputs.images[0].save(image_path)


def main(*, path_experiment: str, experiment_name: str, clean: bool, bucket_name: Optional[str]) -> None:
    tic_total = time.perf_counter()
    start_date = dt.datetime.now(tz=dt.timezone.utc).replace(microsecond=0).isoformat()

    peft_branch = get_peft_branch()
    if peft_branch == "main":
        print_verbose("===== This experiment is categorized as a MAIN run because the PEFT branch is 'main' ======")
    else:
        print_verbose(
            f"===== This experiment is categorized as a TEST run because the PEFT branch is '{peft_branch}' ======"
        )

    peft_config: Optional[PeftConfig] = None
    if os.path.exists(os.path.join(path_experiment, CONFIG_NAME)):
        peft_config = PeftConfig.from_pretrained(path_experiment)
    else:
        print_verbose(f"Could not find PEFT config at {path_experiment}, performing FULL FINETUNING")

    path_train_config = os.path.join(path_experiment, FILE_NAME_TRAIN_PARAMS)
    train_config = get_train_config(path_train_config)
    accelerator_memory_init = init_accelerator()
    set_seed(train_config.seed)

    device_type = infer_device()
    model_info = get_base_model_info(train_config.model_id)
    dataset_info = get_dataset_info(train_config.dataset_id)
    pipeline = get_pipeline(
        model_id=train_config.model_id,
        dtype=train_config.dtype,
        compile=train_config.compile,
        peft_config=peft_config,
        autocast_adapter_dtype=train_config.autocast_adapter_dtype,
        use_gc=train_config.use_gc,
        device_type=device_type,
    )
    print_verbose(pipeline.transformer)

    train_result, prompt_cache = train(
        pipeline=pipeline,
        train_config=train_config,
        accelerator_memory_init=accelerator_memory_init,
        is_adalora=peft_config is not None and peft_config.peft_type == "ADALORA",
        is_adamss=peft_config is not None and peft_config.peft_type == "ADAMSS",
        print_verbose=print_verbose,
        device_type=device_type,
    )

    if train_result.status == TrainStatus.FAILED:
        print_verbose("Training failed, not logging results")
        sys.exit(1)

    file_size = get_file_size(pipeline.transformer, peft_config=peft_config, clean=clean, print_fn=print_verbose)

    time_total = time.perf_counter() - tic_total
    log_results(
        experiment_name=experiment_name,
        train_result=train_result,
        time_total=time_total,
        file_size=file_size,
        model_info=model_info,
        dataset_info=dataset_info,
        start_date=start_date,
        train_config=train_config,
        peft_config=peft_config,
        print_fn=print_verbose,
    )

    if (train_result.status == TrainStatus.SUCCESS) and train_config.sample_image_prompts:
        print_verbose("Generating sample images")
        try:
            sample_image_dir = get_sample_image_save_dir(train_status=train_result.status, peft_branch=peft_branch)
            file_stem = get_artifact_stem(experiment_name, start_date, sample_image_dir)
            generate_sample_images(
                pipeline=pipeline,
                train_config=train_config,
                prompt_cache=prompt_cache,
                sample_image_dir=sample_image_dir,
                file_stem=file_stem,
            )
            print_verbose(f"Stored sample images in {sample_image_dir}")
        except Exception as exc:
            print_verbose(f"Sample image generation failed: {exc}")

    if bucket_name:
        huggingface_hub.create_bucket(bucket_name, exist_ok=True)
        upload_checkpoint_to_bucket(pipeline.transformer, experiment_name, bucket_name)
        upload_images_to_bucket(bucket_name)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("-v", "--verbose", action="store_true", help="Enable verbose output")
    parser.add_argument("path_experiment", type=str, help="Path to the experiment directory")
    parser.add_argument(
        "--clean",
        action="store_true",
        help="Delete training artifacts after run finishes (logs are still saved)",
    )
    parser.add_argument("--bucket_name", type=str, help="HF bucket to upload checkpoints and images to.")
    args = parser.parse_args()

    experiment_name = validate_experiment_path(args.path_experiment)

    if args.verbose:

        def print_verbose(*args, **kwargs) -> None:
            kwargs["file"] = sys.stderr
            print(*args, **kwargs)
    else:

        def print_verbose(*args, **kwargs) -> None:
            pass

    main(
        path_experiment=args.path_experiment,
        experiment_name=experiment_name,
        clean=args.clean,
        bucket_name=args.bucket_name,
    )