From 5d9a5dac92954c5e0a73933971cfd838a445bafd Mon Sep 17 00:00:00 2001 From: scxue Date: Wed, 7 May 2025 16:16:42 +0800 Subject: [PATCH 01/10] test permission From 1123ee7fd7cec8acc7fb697b5f1c4e098a89a9d0 Mon Sep 17 00:00:00 2001 From: scxue Date: Wed, 7 May 2025 16:27:10 +0800 Subject: [PATCH 02/10] Add cross attention type for Sana-Sprint. --- .../models/transformers/sana_transformer.py | 97 ++++++++++++++++++- 1 file changed, 96 insertions(+), 1 deletion(-) diff --git a/src/diffusers/models/transformers/sana_transformer.py b/src/diffusers/models/transformers/sana_transformer.py index 52236275dc88..1d83283882a4 100644 --- a/src/diffusers/models/transformers/sana_transformer.py +++ b/src/diffusers/models/transformers/sana_transformer.py @@ -14,6 +14,7 @@ from typing import Any, Dict, Optional, Tuple, Union +import math import torch import torch.nn.functional as F from torch import nn @@ -184,6 +185,91 @@ def __call__( return hidden_states + +class SanaAttnProcessor3_0: + r""" + Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0). + """ + + def __init__(self): + if not hasattr(F, "scaled_dot_product_attention"): + raise ImportError("SanaAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.") + + @staticmethod + def scaled_dot_product_attention(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, scale=None + ) -> torch.Tensor: + B, H, L, S = *query.size()[:-1], key.size(-2) + scale_factor = 1 / math.sqrt(query.size(-1)) if scale is None else scale + attn_bias = torch.zeros(B, H, L, S, dtype=query.dtype, device=query.device) + + if attn_mask is not None: + if attn_mask.dtype == torch.bool: + attn_bias.masked_fill_(attn_mask.logical_not(), float("-inf")) + else: + attn_bias += attn_mask + attn_weight = query @ key.transpose(-2, -1) * scale_factor + attn_weight += attn_bias + attn_weight = torch.softmax(attn_weight, dim=-1) + attn_weight = torch.dropout(attn_weight, dropout_p, train=True) + return attn_weight @ value + + # return x + def __call__( + self, + attn: Attention, + hidden_states: torch.Tensor, + encoder_hidden_states: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + batch_size, sequence_length, _ = ( + hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape + ) + + if attention_mask is not None: + attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size) + # scaled_dot_product_attention expects attention_mask shape to be + # (batch, heads, source_length, target_length) + attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1]) + + query = attn.to_q(hidden_states) + + if encoder_hidden_states is None: + encoder_hidden_states = hidden_states + + key = attn.to_k(encoder_hidden_states) + value = attn.to_v(encoder_hidden_states) + + if attn.norm_q is not None: + query = attn.norm_q(query) + if attn.norm_k is not None: + key = attn.norm_k(key) + + inner_dim = key.shape[-1] + head_dim = inner_dim // attn.heads + + query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + + key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + + # the output of sdp = (batch, num_heads, seq_len, head_dim) + # TODO: add support for attn.scale when we move to Torch 2.1 + hidden_states = self.scaled_dot_product_attention( + query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False + ) + + hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) + hidden_states = hidden_states.to(query.dtype) + + # linear proj + hidden_states = attn.to_out[0](hidden_states) + # dropout + hidden_states = attn.to_out[1](hidden_states) + + hidden_states = hidden_states / attn.rescale_output_factor + + return hidden_states + class SanaTransformerBlock(nn.Module): r""" @@ -205,6 +291,7 @@ def __init__( attention_out_bias: bool = True, mlp_ratio: float = 2.5, qk_norm: Optional[str] = None, + cross_attention_type: str = "flash", ) -> None: super().__init__() @@ -223,6 +310,12 @@ def __init__( ) # 2. Cross Attention + if cross_attention_type == "flash": + cross_attention_processor = SanaAttnProcessor2_0() + elif cross_attention_type == "vanilla": + cross_attention_processor = SanaAttnProcessor3_0() + else: + raise ValueError(f"Cross attention type {cross_attention_type} is not defined.") if cross_attention_dim is not None: self.norm2 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps) self.attn2 = Attention( @@ -235,7 +328,7 @@ def __init__( dropout=dropout, bias=True, out_bias=attention_out_bias, - processor=SanaAttnProcessor2_0(), + processor=cross_attention_processor, ) # 3. Feed-forward @@ -360,6 +453,7 @@ def __init__( guidance_embeds_scale: float = 0.1, qk_norm: Optional[str] = None, timestep_scale: float = 1.0, + cross_attention_type: str = "flash", ) -> None: super().__init__() @@ -402,6 +496,7 @@ def __init__( norm_eps=norm_eps, mlp_ratio=mlp_ratio, qk_norm=qk_norm, + cross_attention_type=cross_attention_type, ) for _ in range(num_layers) ] From acefec8cf8ba3d3249973c1dacba68beadf803bd Mon Sep 17 00:00:00 2001 From: scxue Date: Wed, 7 May 2025 17:33:42 +0800 Subject: [PATCH 03/10] Add Sana-Sprint training script in diffusers. --- .../sana/train_sana_sprint_diffusers.py | 1680 +++++++++++++++++ .../sana/train_sana_sprint_diffusers.sh | 24 + 2 files changed, 1704 insertions(+) create mode 100644 examples/research_projects/sana/train_sana_sprint_diffusers.py create mode 100644 examples/research_projects/sana/train_sana_sprint_diffusers.sh diff --git a/examples/research_projects/sana/train_sana_sprint_diffusers.py b/examples/research_projects/sana/train_sana_sprint_diffusers.py new file mode 100644 index 000000000000..15aa323953b1 --- /dev/null +++ b/examples/research_projects/sana/train_sana_sprint_diffusers.py @@ -0,0 +1,1680 @@ +#!/usr/bin/env python +# coding=utf-8 +# Copyright 2025 Sana-Sprint team. All rights reserved. +# Copyright 2025 The HuggingFace Inc. team. All rights reserved. +# +# 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 + +import argparse +import copy +import itertools +import json +import logging +import math +import os +import random +import shutil +import warnings +from pathlib import Path + +import accelerate +import io +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.checkpoint +import torchvision.transforms as T +import torchvision.transforms.functional as TF +import transformers +import webdataset as wds +from accelerate import Accelerator +from accelerate.logging import get_logger +from accelerate.utils import DistributedDataParallelKwargs, DistributedType, ProjectConfiguration, set_seed +from braceexpand import braceexpand +from datasets import load_dataset +from huggingface_hub import create_repo, upload_folder +from huggingface_hub.utils import insecure_hashlib +from packaging import version +from peft.utils import get_peft_model_state_dict +from PIL import Image +from PIL.ImageOps import exif_transpose +from safetensors.torch import load_file +from torch.nn.utils.spectral_norm import SpectralNorm +from torch.utils.data import default_collate, Dataset, DataLoader +from torchvision.transforms.functional import crop +from tqdm.auto import tqdm +from transformers import AutoTokenizer, Gemma2Model +from typing import Callable, List, Union +from webdataset.tariterators import ( + base_plus_ext, + tar_file_expander, + url_opener, + valid_sample, +) + +import diffusers +from diffusers import ( + AutoencoderDC, + FlowMatchEulerDiscreteScheduler, + SanaPipeline, + SanaSprintPipeline, + SanaTransformer2DModel, + SCMScheduler, +) +from diffusers.optimization import get_scheduler +from diffusers.training_utils import ( + cast_training_params, + compute_density_for_timestep_sampling, + compute_loss_weighting_for_sd3, + free_memory, +) +from diffusers.utils import ( + check_min_version, + convert_unet_state_dict_to_peft, + is_wandb_available, +) +from diffusers.utils.hub_utils import load_or_create_model_card, populate_model_card +from diffusers.utils.import_utils import is_torch_npu_available +from diffusers.utils.torch_utils import is_compiled_module + + +if is_wandb_available(): + import wandb + +# Will error if the minimal version of diffusers is not installed. Remove at your own risks. +check_min_version("0.33.0.dev0") + +logger = get_logger(__name__) + +if is_torch_npu_available(): + torch.npu.config.allow_internal_format = False + +COMPLEX_HUMAN_INSTRUCTION = [ + "Given a user prompt, generate an 'Enhanced prompt' that provides detailed visual descriptions suitable for image generation. Evaluate the level of detail in the user prompt:", + "- If the prompt is simple, focus on adding specifics about colors, shapes, sizes, textures, and spatial relationships to create vivid and concrete scenes.", + "- If the prompt is already detailed, refine and enhance the existing details slightly without overcomplicating.", + "Here are examples of how to transform or refine prompts:", + "- User Prompt: A cat sleeping -> Enhanced: A small, fluffy white cat curled up in a round shape, sleeping peacefully on a warm sunny windowsill, surrounded by pots of blooming red flowers.", + "- User Prompt: A busy city street -> Enhanced: A bustling city street scene at dusk, featuring glowing street lamps, a diverse crowd of people in colorful clothing, and a double-decker bus passing by towering glass skyscrapers.", + "Please generate only the enhanced description for the prompt below and avoid including any additional commentary or evaluations:", + "User Prompt: ", +] + + + +class Text2ImageDataset(Dataset): + """ + A PyTorch Dataset class for loading text-image pairs from a HuggingFace dataset. + This dataset is designed for text-to-image generation tasks. + + Args: + hf_dataset (datasets.Dataset): A HuggingFace dataset containing 'image' (bytes) + and 'llava' (text) fields. Note that 'llava' is the field name for text descriptions in this specific dataset - you may need to adjust this key if using a different HuggingFace dataset with a different text field name. + resolution (int, optional): Target resolution for image resizing. Defaults to 1024. + + Returns: + dict: A dictionary containing: + - 'text': The text description (str) + - 'image': The processed image tensor (torch.Tensor) of shape [3, resolution, resolution] + """ + def __init__(self, hf_dataset, resolution=1024): + self.dataset = hf_dataset + self.transform = T.Compose([ + T.Lambda(lambda img: img.convert("RGB")), + T.Resize(resolution), # Image.BICUBIC + T.CenterCrop(resolution), + # T.RandomHorizontalFlip(), + T.ToTensor(), + T.Normalize([0.5], [0.5]), + ]) + + def __len__(self): + return len(self.dataset) + + def __getitem__(self, idx): + item = self.dataset[idx] + text = item['llava'] + image_bytes = item['image'] + + # Convert bytes to PIL Image + image = Image.open(io.BytesIO(image_bytes)) + + image_tensor = self.transform(image) + + return { + 'text': text, + 'image': image_tensor + } + +# TODO here +def save_model_card( + repo_id: str, + images=None, + base_model: str = None, + validation_prompt=None, + repo_folder=None, +): + widget_dict = [] + if images is not None: + for i, image in enumerate(images): + image.save(os.path.join(repo_folder, f"image_{i}.png")) + widget_dict.append( + {"text": validation_prompt if validation_prompt else " ", "output": {"url": f"image_{i}.png"}} + ) + + model_description = f""" +# Sana Sprint - {repo_id} + + + +## Model description + +These are {repo_id} Sana Sprint weights for {base_model}. + +The weights were trained using [Sana-Sprint](https://nvlabs.github.io/Sana/Sprint/). + +## License + +TODO +""" + model_card = load_or_create_model_card( + repo_id_or_path=repo_id, + from_training=True, + license="other", + base_model=base_model, + model_description=model_description, + widget=widget_dict, + ) + tags = [ + "text-to-image", + "diffusers-training", + "diffusers", + "sana-sprint", + "sana-sprint-diffusers", + ] + + model_card = populate_model_card(model_card, tags=tags) + model_card.save(os.path.join(repo_folder, "README.md")) + + +def log_validation( + pipeline, + args, + accelerator, + pipeline_args, + epoch, + is_final_validation=False, +): + logger.info( + f"Running validation... \n Generating {args.num_validation_images} images with prompt:" + f" {args.validation_prompt}." + ) + if args.enable_vae_tiling: + pipeline.vae.enable_tiling(tile_sample_min_height=1024, tile_sample_stride_width=1024) + + pipeline.text_encoder = pipeline.text_encoder.to(torch.bfloat16) + pipeline = pipeline.to(accelerator.device) + pipeline.set_progress_bar_config(disable=True) + + # run inference + generator = torch.Generator(device=accelerator.device).manual_seed(args.seed) if args.seed is not None else None + + images = [pipeline(**pipeline_args, generator=generator).images[0] for _ in range(args.num_validation_images)] + + for tracker in accelerator.trackers: + phase_name = "test" if is_final_validation else "validation" + if tracker.name == "tensorboard": + np_images = np.stack([np.asarray(img) for img in images]) + tracker.writer.add_images(phase_name, np_images, epoch, dataformats="NHWC") + if tracker.name == "wandb": + tracker.log( + { + phase_name: [ + wandb.Image(image, caption=f"{i}: {args.validation_prompt}") for i, image in enumerate(images) + ] + } + ) + + del pipeline + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + return images + + +def parse_args(input_args=None): + parser = argparse.ArgumentParser(description="Simple example of a training script.") + parser.add_argument( + "--pretrained_model_name_or_path", + type=str, + default=None, + required=True, + help="Path to pretrained model or model identifier from huggingface.co/models.", + ) + parser.add_argument( + "--revision", + type=str, + default=None, + required=False, + help="Revision of pretrained model identifier from huggingface.co/models.", + ) + parser.add_argument( + "--variant", + type=str, + default=None, + help="Variant of the model files of the pretrained model identifier from huggingface.co/models, 'e.g.' fp16", + ) + + parser.add_argument( + "--cache_dir", + type=str, + default=None, + help="The directory where the downloaded models and datasets will be stored.", + ) + + parser.add_argument( + "--image_column", + type=str, + default="image", + help="The column of the dataset containing the target image. By " + "default, the standard Image Dataset maps out 'file_name' " + "to 'image'.", + ) + parser.add_argument( + "--caption_column", + type=str, + default=None, + help="The column of the dataset containing the instance prompt for each image", + ) + + parser.add_argument("--repeats", type=int, default=1, help="How many times to repeat the training data.") + parser.add_argument( + "--max_sequence_length", + type=int, + default=300, + help="Maximum sequence length to use with with the Gemma model", + ) + parser.add_argument( + "--validation_prompt", + type=str, + default=None, + help="A prompt that is used during validation to verify that the model is learning.", + ) + parser.add_argument( + "--num_validation_images", + type=int, + default=4, + help="Number of images that should be generated during validation with `validation_prompt`.", + ) + parser.add_argument( + "--validation_epochs", + type=int, + default=50, + help=( + "Run dreambooth validation every X epochs. Dreambooth validation consists of running the prompt" + " `args.validation_prompt` multiple times: `args.num_validation_images`." + ), + ) + parser.add_argument( + "--output_dir", + type=str, + default="sana-dreambooth-lora", + help="The output directory where the model predictions and checkpoints will be written.", + ) + parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.") + # ----Image Processing---- + parser.add_argument("--file_path", nargs='+', required=True, + help="List of parquet files (space-separated)") + parser.add_argument( + "--dataset_name", + type=str, + default=None, + help=( + "The name of the Dataset (from the HuggingFace hub) to train on (could be your own, possibly private," + " dataset). It can also be a path pointing to a local copy of a dataset in your filesystem," + " or to a folder containing files that 🤗 Datasets can understand." + ), + ) + parser.add_argument( + "--resolution", + type=int, + default=512, + help=( + "The resolution for input images, all the images in the train/validation dataset will be resized to this" + " resolution" + ), + ) + parser.add_argument( + "--use_fix_crop_and_size", + action="store_true", + help="Whether or not to use the fixed crop and size for the teacher model.", + default=False, + ) + parser.add_argument( + "--center_crop", + default=False, + action="store_true", + help=( + "Whether to center crop the input images to the resolution. If not set, the images will be randomly" + " cropped. The images will be resized to the resolution first before cropping." + ), + ) + parser.add_argument( + "--random_flip", + action="store_true", + help="whether to randomly flip images horizontally", + ) + parser.add_argument( + "--train_batch_size", type=int, default=4, help="Batch size (per device) for the training dataloader." + ) + parser.add_argument( + "--sample_batch_size", type=int, default=4, help="Batch size (per device) for sampling images." + ) + parser.add_argument("--num_train_epochs", type=int, default=1) + parser.add_argument( + "--max_train_steps", + type=int, + default=None, + help="Total number of training steps to perform. If provided, overrides num_train_epochs.", + ) + parser.add_argument( + "--checkpointing_steps", + type=int, + default=500, + help=( + "Save a checkpoint of the training state every X updates. These checkpoints can be used both as final" + " checkpoints in case they are better than the last checkpoint, and are also suitable for resuming" + " training using `--resume_from_checkpoint`." + ), + ) + parser.add_argument( + "--checkpoints_total_limit", + type=int, + default=None, + help=("Max number of checkpoints to store."), + ) + parser.add_argument( + "--resume_from_checkpoint", + type=str, + default=None, + help=( + "Whether training should be resumed from a previous checkpoint. Use a path saved by" + ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.' + ), + ) + parser.add_argument( + "--gradient_accumulation_steps", + type=int, + default=1, + help="Number of updates steps to accumulate before performing a backward/update pass.", + ) + parser.add_argument( + "--gradient_checkpointing", + action="store_true", + help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.", + ) + parser.add_argument( + "--learning_rate", + type=float, + default=1e-4, + help="Initial learning rate (after the potential warmup period) to use.", + ) + parser.add_argument( + "--scale_lr", + action="store_true", + default=False, + help="Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.", + ) + parser.add_argument( + "--lr_scheduler", + type=str, + default="constant", + help=( + 'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",' + ' "constant", "constant_with_warmup"]' + ), + ) + parser.add_argument( + "--lr_warmup_steps", type=int, default=500, help="Number of steps for the warmup in the lr scheduler." + ) + parser.add_argument( + "--lr_num_cycles", + type=int, + default=1, + help="Number of hard resets of the lr in cosine_with_restarts scheduler.", + ) + parser.add_argument("--lr_power", type=float, default=1.0, help="Power factor of the polynomial scheduler.") + parser.add_argument( + "--dataloader_num_workers", + type=int, + default=0, + help=( + "Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process." + ), + ) + parser.add_argument( + "--logit_mean", type=float, default=0.2, help="mean to use when using the `'logit_normal'` weighting scheme." + ) + parser.add_argument( + "--logit_std", type=float, default=1.6, help="std to use when using the `'logit_normal'` weighting scheme." + ) + parser.add_argument("--logit_mean_discriminator", type=float, default=-0.6, help="Logit mean for discriminator timestep sampling") + parser.add_argument("--logit_std_discriminator", type=float, default=1.0, help="Logit std for discriminator timestep sampling") + parser.add_argument("--ladd_multi_scale", action="store_true", help="Whether to use multi-scale discriminator") + parser.add_argument("--head_block_ids", type=int, nargs="+", default=[2, 8, 14, 19], help="Specify which transformer blocks to use for discriminator heads") + parser.add_argument("--adv_lambda", type=float, default=0.5, help="Weighting coefficient for adversarial loss") + parser.add_argument("--scm_lambda", type=float, default=1.0, help="Weighting coefficient for SCM loss") + parser.add_argument("--gradient_clip", type=float, default=0.1, help="Threshold for gradient clipping") + parser.add_argument("--sigma_data", type=float, default=0.5, help="Standard deviation of data distribution is supposed to be 0.5") + parser.add_argument("--tangent_warmup_steps", type=int, default=4000, help="Number of warmup steps for tangent vectors") + parser.add_argument("--guidance_embeds_scale", type=float, default=0.1, help="Scaling factor for guidance embeddings") + parser.add_argument("--scm_cfg_scale", type=float, nargs="+", default=[4, 4.5, 5], help="Range for classifier-free guidance scale") + parser.add_argument("--train_largest_timestep", action="store_true", help="Whether to enable special training for large timesteps") + parser.add_argument("--largest_timestep", type=float, default=1.57080, help="Maximum timestep value") + parser.add_argument("--largest_timestep_prob", type=float, default=0.5, help="Sampling probability for large timesteps") + parser.add_argument("--misaligned_pairs_D", action="store_true", help="Add misaligned sample pairs for discriminator") + parser.add_argument( + "--optimizer", + type=str, + default="AdamW", + help=('The optimizer type to use. Choose between ["AdamW", "prodigy"]'), + ) + + parser.add_argument( + "--use_8bit_adam", + action="store_true", + help="Whether or not to use 8-bit Adam from bitsandbytes. Ignored if optimizer is not set to AdamW", + ) + + parser.add_argument( + "--adam_beta1", type=float, default=0.9, help="The beta1 parameter for the Adam and Prodigy optimizers." + ) + parser.add_argument( + "--adam_beta2", type=float, default=0.999, help="The beta2 parameter for the Adam and Prodigy optimizers." + ) + parser.add_argument( + "--prodigy_beta3", + type=float, + default=None, + help="coefficients for computing the Prodigy stepsize using running averages. If set to None, " + "uses the value of square root of beta2. Ignored if optimizer is adamW", + ) + parser.add_argument("--prodigy_decouple", type=bool, default=True, help="Use AdamW style decoupled weight decay") + parser.add_argument("--adam_weight_decay", type=float, default=1e-04, help="Weight decay to use for unet params") + + parser.add_argument( + "--adam_epsilon", + type=float, + default=1e-08, + help="Epsilon value for the Adam optimizer and Prodigy optimizers.", + ) + + parser.add_argument( + "--prodigy_use_bias_correction", + type=bool, + default=True, + help="Turn on Adam's bias correction. True by default. Ignored if optimizer is adamW", + ) + parser.add_argument( + "--prodigy_safeguard_warmup", + type=bool, + default=True, + help="Remove lr from the denominator of D estimate to avoid issues during warm-up stage. True by default. " + "Ignored if optimizer is adamW", + ) + parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.") + parser.add_argument("--push_to_hub", action="store_true", help="Whether or not to push the model to the Hub.") + parser.add_argument("--hub_token", type=str, default=None, help="The token to use to push to the Model Hub.") + parser.add_argument( + "--hub_model_id", + type=str, + default=None, + help="The name of the repository to keep in sync with the local `output_dir`.", + ) + parser.add_argument( + "--logging_dir", + type=str, + default="logs", + help=( + "[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to" + " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***." + ), + ) + parser.add_argument( + "--allow_tf32", + action="store_true", + help=( + "Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see" + " https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices" + ), + ) + parser.add_argument( + "--cache_latents", + action="store_true", + default=False, + help="Cache the VAE latents", + ) + parser.add_argument( + "--report_to", + type=str, + default="tensorboard", + help=( + 'The integration to report the results and logs to. Supported platforms are `"tensorboard"`' + ' (default), `"wandb"` and `"comet_ml"`. Use `"all"` to report to all integrations.' + ), + ) + parser.add_argument( + "--mixed_precision", + type=str, + default=None, + 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. Default to the value of accelerate config of the current system or the" + " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config." + ), + ) + parser.add_argument( + "--upcast_before_saving", + action="store_true", + default=False, + help=( + "Whether to upcast the trained transformer layers to float32 before saving (at the end of training). " + "Defaults to precision dtype used for training to save memory" + ), + ) + parser.add_argument( + "--offload", + action="store_true", + help="Whether to offload the VAE and the text encoder to CPU when they are not used.", + ) + parser.add_argument("--local_rank", type=int, default=-1, help="For distributed training: local_rank") + parser.add_argument("--enable_vae_tiling", action="store_true", help="Enabla vae tiling in log validation") + parser.add_argument("--enable_npu_flash_attention", action="store_true", help="Enabla Flash Attention for NPU") + + if input_args is not None: + args = parser.parse_args(input_args) + else: + 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 + + + return args + + +class ResidualBlock(nn.Module): + def __init__(self, fn: Callable): + super().__init__() + self.fn = fn + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return (self.fn(x) + x) / np.sqrt(2) + + +class SpectralConv1d(nn.Conv1d): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + SpectralNorm.apply(self, name="weight", n_power_iterations=1, dim=0, eps=1e-12) + + +class BatchNormLocal(nn.Module): + def __init__(self, num_features: int, affine: bool = True, virtual_bs: int = 8, eps: float = 1e-5): + super().__init__() + self.virtual_bs = virtual_bs + self.eps = eps + self.affine = affine + + if self.affine: + self.weight = nn.Parameter(torch.ones(num_features)) + self.bias = nn.Parameter(torch.zeros(num_features)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + shape = x.size() + + # Reshape batch into groups. + G = np.ceil(x.size(0) / self.virtual_bs).astype(int) + x = x.view(G, -1, x.size(-2), x.size(-1)) + + # Calculate stats. + mean = x.mean([1, 3], keepdim=True) + var = x.var([1, 3], keepdim=True, unbiased=False) + x = (x - mean) / (torch.sqrt(var + self.eps)) + + if self.affine: + x = x * self.weight[None, :, None] + self.bias[None, :, None] + + return x.view(shape) + + +def make_block(channels: int, kernel_size: int) -> nn.Module: + return nn.Sequential( + SpectralConv1d( + channels, + channels, + kernel_size=kernel_size, + padding=kernel_size // 2, + padding_mode="circular", + ), + BatchNormLocal(channels), + nn.LeakyReLU(0.2, True), + ) + + +# Adapted from https://github.com/autonomousvision/stylegan-t/blob/main/networks/discriminator.py +class DiscHead(nn.Module): + def __init__(self, channels: int, c_dim: int, cmap_dim: int = 64): + super().__init__() + self.channels = channels + self.c_dim = c_dim + self.cmap_dim = cmap_dim + + self.main = nn.Sequential( + make_block(channels, kernel_size=1), ResidualBlock(make_block(channels, kernel_size=9)) + ) + + if self.c_dim > 0: + self.cmapper = nn.Linear(self.c_dim, cmap_dim) + self.cls = SpectralConv1d(channels, cmap_dim, kernel_size=1, padding=0) + else: + self.cls = SpectralConv1d(channels, 1, kernel_size=1, padding=0) + + def forward(self, x: torch.Tensor, c: torch.Tensor) -> torch.Tensor: + h = self.main(x) + out = self.cls(h) + + if self.c_dim > 0: + cmap = self.cmapper(c).unsqueeze(-1) + out = (out * cmap).sum(1, keepdim=True) * (1 / np.sqrt(self.cmap_dim)) + + return out + + +class SanaMSCMDiscriminator(nn.Module): + def __init__(self, pretrained_model, is_multiscale=False, head_block_ids=None): + super().__init__() + self.transformer = pretrained_model + self.transformer.requires_grad_(False) + + if head_block_ids is None or len(head_block_ids) == 0: + self.block_hooks = {2, 8, 14, 20, 27} if is_multiscale else {self.transformer.depth - 1} + else: + self.block_hooks = head_block_ids + + heads = [] + for i in range(len(self.block_hooks)): + heads.append(DiscHead(self.transformer.hidden_size, 0, 0)) + self.heads = nn.ModuleList(heads) + + def get_head_inputs(self): + return self.head_inputs + + def forward(self, hidden_states, timestep, encoder_hidden_states=None, **kwargs): + feat_list = [] + self.head_inputs = [] + + def get_features(module, input, output): + feat_list.append(output) + return output + + hooks = [] + for i, block in enumerate(self.transformer.transformer_blocks): + if i in self.block_hooks: + hooks.append(block.register_forward_hook(get_features)) + + self.transformer(hidden_states=hidden_states, timestep=timestep, encoder_hidden_states=encoder_hidden_states, return_logvar=False, **kwargs) + + for hook in hooks: + hook.remove() + + res_list = [] + for feat, head in zip(feat_list, self.heads): + B, N, C = feat.shape + feat = feat.transpose(1, 2) # [B, C, N] + self.head_inputs.append(feat) + res_list.append(head(feat, None).reshape(feat.shape[0], -1)) + + concat_res = torch.cat(res_list, dim=1) + + return concat_res + + @property + def model(self): + return self.transformer + + def save_pretrained(self, path): + torch.save(self.state_dict(), path) + + +class DiscHeadModel: + def __init__(self, disc): + self.disc = disc + + def state_dict(self): + return {name: param for name, param in self.disc.state_dict().items() if not name.startswith("transformer.")} + + def __getattr__(self, name): + return getattr(self.disc, name) + +class SanaTrigFlow(SanaTransformer2DModel): + def __init__(self, original_model, guidance=False): + self.__dict__ = original_model.__dict__ + self.hidden_size = self.config.num_attention_heads * self.config.attention_head_dim + self.guidance = guidance + if self.guidance: + hidden_size = self.config.num_attention_heads * self.config.attention_head_dim + self.logvar_linear = torch.nn.Linear(hidden_size, 1) + torch.nn.init.xavier_uniform_(self.logvar_linear.weight) + torch.nn.init.constant_(self.logvar_linear.bias, 0) + + def forward(self, hidden_states, encoder_hidden_states, timestep, guidance=None, jvp=False, return_logvar=False, **kwargs): + batch_size = hidden_states.shape[0] + latents = hidden_states + prompt_embeds = encoder_hidden_states + t = timestep + + # TrigFlow --> Flow Transformation + timestep = t.expand(latents.shape[0]).to(prompt_embeds.dtype) + latents_model_input = latents + + flow_timestep = torch.sin(timestep) / (torch.cos(timestep) + torch.sin(timestep)) + + flow_timestep_expanded = flow_timestep.view(-1, 1, 1, 1) + latent_model_input = latents_model_input * torch.sqrt( + flow_timestep_expanded**2 + (1 - flow_timestep_expanded) ** 2 + ) + latent_model_input = latent_model_input.to(prompt_embeds.dtype) + + # forward in original flow + + if jvp and self.gradient_checkpointing: + self.gradient_checkpointing = False + model_out = super().forward(hidden_states=latent_model_input, encoder_hidden_states=prompt_embeds, timestep=flow_timestep, guidance=guidance, **kwargs)[0] + self.gradient_checkpointing = True + else: + model_out = super().forward(hidden_states=latent_model_input, encoder_hidden_states=prompt_embeds, timestep=flow_timestep, guidance=guidance, **kwargs)[0] + + + # Flow --> TrigFlow Transformation + trigflow_model_out = ((1 - 2 * flow_timestep_expanded) * latent_model_input + (1 - 2 * flow_timestep_expanded + 2 * flow_timestep_expanded**2) * model_out) / torch.sqrt( + flow_timestep_expanded**2 + (1 - flow_timestep_expanded) ** 2 + ) + + + if self.guidance and guidance is not None: + timestep, embedded_timestep = self.time_embed( + timestep, guidance=guidance, hidden_dtype=hidden_states.dtype + ) + else: + timestep, embedded_timestep = self.time_embed( + timestep, batch_size=batch_size, hidden_dtype=hidden_states.dtype + ) + + if return_logvar: + logvar = self.logvar_linear(embedded_timestep) + return trigflow_model_out, logvar + + + return (trigflow_model_out,) + + + +def compute_density_for_timestep_sampling_scm( + batch_size: int, logit_mean: float = None, logit_std: float = None +): + """Compute the density for sampling the timesteps when doing Sana-Sprint training. + """ + sigma = torch.randn(batch_size, device="cpu") + sigma = (sigma * logit_std + logit_mean).exp() + u = torch.atan(sigma / 0.5) # TODO: 0.5 should be a hyper-parameter + + return u + + + +def main(args): + if args.report_to == "wandb" and args.hub_token is not None: + raise ValueError( + "You cannot use both --report_to=wandb and --hub_token due to a security risk of exposing your token." + " Please use `huggingface-cli login` to authenticate with the Hub." + ) + + if torch.backends.mps.is_available() and args.mixed_precision == "bf16": + # due to pytorch#99272, MPS does not yet support bfloat16. + raise ValueError( + "Mixed precision training with bfloat16 is not supported on MPS. Please use fp16 (recommended) or fp32 instead." + ) + + logging_dir = Path(args.output_dir, args.logging_dir) + + accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir) + kwargs = DistributedDataParallelKwargs(find_unused_parameters=True) + accelerator = Accelerator( + gradient_accumulation_steps=args.gradient_accumulation_steps, + mixed_precision=args.mixed_precision, + log_with=args.report_to, + project_config=accelerator_project_config, + kwargs_handlers=[kwargs], + ) + + # Disable AMP for MPS. + if torch.backends.mps.is_available(): + accelerator.native_amp = False + + if args.report_to == "wandb": + if not is_wandb_available(): + raise ImportError("Make sure to install wandb if you want to use it for logging during training.") + + # Make one log on every process with the configuration for debugging. + logging.basicConfig( + format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", + datefmt="%m/%d/%Y %H:%M:%S", + level=logging.INFO, + ) + logger.info(accelerator.state, main_process_only=False) + if accelerator.is_local_main_process: + transformers.utils.logging.set_verbosity_warning() + diffusers.utils.logging.set_verbosity_info() + else: + transformers.utils.logging.set_verbosity_error() + diffusers.utils.logging.set_verbosity_error() + + # If passed along, set the training seed now. + if args.seed is not None: + set_seed(args.seed) + + + # Handle the repository creation + if accelerator.is_main_process: + if args.output_dir is not None: + os.makedirs(args.output_dir, exist_ok=True) + + if args.push_to_hub: + repo_id = create_repo( + repo_id=args.hub_model_id or Path(args.output_dir).name, + exist_ok=True, + ).repo_id + + # Load the tokenizer + tokenizer = AutoTokenizer.from_pretrained( + args.pretrained_model_name_or_path, + subfolder="tokenizer", + revision=args.revision, + ) + + # Load scheduler and models + text_encoder = Gemma2Model.from_pretrained( + args.pretrained_model_name_or_path, subfolder="text_encoder", revision=args.revision, variant=args.variant + ) + vae = AutoencoderDC.from_pretrained( + args.pretrained_model_name_or_path, + subfolder="vae", + revision=args.revision, + variant=args.variant, + ) + + ori_transformer = SanaTransformer2DModel.from_pretrained( + args.pretrained_model_name_or_path, subfolder="transformer", revision=args.revision, variant=args.variant, + guidance_embeds=True, cross_attention_type='vanilla' + ) + + ori_transformer_no_guide = SanaTransformer2DModel.from_pretrained( + args.pretrained_model_name_or_path, subfolder="transformer", revision=args.revision, variant=args.variant, + guidance_embeds=False + ) + + original_state_dict = load_file(f"{args.pretrained_model_name_or_path}/transformer/diffusion_pytorch_model.safetensors") + + param_mapping = { + 'time_embed.emb.timestep_embedder.linear_1.weight': 'time_embed.timestep_embedder.linear_1.weight', + 'time_embed.emb.timestep_embedder.linear_1.bias': 'time_embed.timestep_embedder.linear_1.bias', + 'time_embed.emb.timestep_embedder.linear_2.weight': 'time_embed.timestep_embedder.linear_2.weight', + 'time_embed.emb.timestep_embedder.linear_2.bias': 'time_embed.timestep_embedder.linear_2.bias', + } + + for src_key, dst_key in param_mapping.items(): + if src_key in original_state_dict: + ori_transformer.load_state_dict({dst_key: original_state_dict[src_key]}, strict=False, assign=True) + + guidance_embedder_module = ori_transformer.time_embed.guidance_embedder + + zero_state_dict = {} + + + target_device = accelerator.device + param_w1 = guidance_embedder_module.linear_1.weight + zero_state_dict['linear_1.weight'] = torch.zeros(param_w1.shape, device=target_device) + param_b1 = guidance_embedder_module.linear_1.bias + zero_state_dict['linear_1.bias'] = torch.zeros(param_b1.shape, device=target_device) + param_w2 = guidance_embedder_module.linear_2.weight + zero_state_dict['linear_2.weight'] = torch.zeros(param_w2.shape, device=target_device) + param_b2 = guidance_embedder_module.linear_2.bias + zero_state_dict['linear_2.bias'] = torch.zeros(param_b2.shape, device=target_device) + guidance_embedder_module.load_state_dict(zero_state_dict, strict=False, assign=True) + + + transformer = SanaTrigFlow(ori_transformer, guidance=True).train() + pretrained_model = SanaTrigFlow(ori_transformer_no_guide, guidance=False).eval() + + disc = SanaMSCMDiscriminator( + pretrained_model, + is_multiscale=args.ladd_multi_scale, + head_block_ids=args.head_block_ids, + ).train() + + + transformer.requires_grad_(True) + pretrained_model.requires_grad_(False) + disc.model.requires_grad_(False) + disc.heads.requires_grad_(True) + vae.requires_grad_(False) + text_encoder.requires_grad_(False) + + # For mixed precision training we cast all non-trainable weights (vae, text_encoder and transformer) to half-precision + # as these weights are only used for inference, keeping weights in full precision is not required. + weight_dtype = torch.float32 + if accelerator.mixed_precision == "fp16": + weight_dtype = torch.float16 + elif accelerator.mixed_precision == "bf16": + weight_dtype = torch.bfloat16 + + if torch.backends.mps.is_available() and weight_dtype == torch.bfloat16: + # due to pytorch#99272, MPS does not yet support bfloat16. + raise ValueError( + "Mixed precision training with bfloat16 is not supported on MPS. Please use fp16 (recommended) or fp32 instead." + ) + + # VAE should always be kept in fp32 for SANA (?) + vae.to(accelerator.device, dtype=torch.float32) + transformer.to(accelerator.device, dtype=weight_dtype) + pretrained_model.to(accelerator.device, dtype=weight_dtype) + disc.to(accelerator.device, dtype=weight_dtype) + # because Gemma2 is particularly suited for bfloat16. + text_encoder.to(dtype=torch.bfloat16) + + if args.enable_npu_flash_attention: + if is_torch_npu_available(): + logger.info("npu flash attention enabled.") + for block in transformer.transformer_blocks: + block.attn2.set_use_npu_flash_attention(True) + for block in pretrained_model.transformer_blocks: + block.attn2.set_use_npu_flash_attention(True) + else: + raise ValueError("npu flash attention requires torch_npu extensions and is supported only on npu device ") + + # Initialize a text encoding pipeline and keep it to CPU for now. + text_encoding_pipeline = SanaPipeline.from_pretrained( + args.pretrained_model_name_or_path, + vae=None, + transformer=None, + text_encoder=text_encoder, + tokenizer=tokenizer, + torch_dtype=torch.bfloat16, + ) + text_encoding_pipeline = text_encoding_pipeline.to(accelerator.device) + + if args.gradient_checkpointing: + transformer.enable_gradient_checkpointing() + + + def unwrap_model(model): + model = accelerator.unwrap_model(model) + model = model._orig_mod if is_compiled_module(model) else model + return model + + if version.parse(accelerate.__version__) >= version.parse("0.16.0"): + + def save_model_hook(models, weights, output_dir): + if accelerator.is_main_process: + for model in models: + unwrapped_model = unwrap_model(model) + # Handle transformer model + if isinstance(unwrapped_model, type(unwrap_model(transformer))): + model = unwrapped_model + model.save_pretrained(os.path.join(output_dir, "transformer")) + # Handle discriminator model (only save heads) + elif isinstance(unwrapped_model, type(unwrap_model(disc))): + # Save only the heads + torch.save( + unwrapped_model.heads.state_dict(), + os.path.join(output_dir, "disc_heads.pt") + ) + else: + raise ValueError(f"unexpected save model: {unwrapped_model.__class__}") + + # make sure to pop weight so that corresponding model is not saved again + if weights: + weights.pop() + + def load_model_hook(models, input_dir): + transformer_ = None + disc_ = None + + if not accelerator.distributed_type == DistributedType.DEEPSPEED: + while len(models) > 0: + model = models.pop() + unwrapped_model = unwrap_model(model) + + if isinstance(unwrapped_model, type(unwrap_model(transformer))): + transformer_ = model # noqa: F841 + elif isinstance(unwrapped_model, type(unwrap_model(disc))): + # Load only the heads + heads_state_dict = torch.load(os.path.join(input_dir, "disc_heads.pt")) + unwrapped_model.heads.load_state_dict(heads_state_dict) + disc_ = model # noqa: F841 + else: + raise ValueError(f"unexpected save model: {unwrapped_model.__class__}") + + else: + # DeepSpeed case + transformer_ = SanaTransformer2DModel.from_pretrained(input_dir, subfolder="transformer") # noqa: F841 + disc_heads_state_dict = torch.load(os.path.join(input_dir, "disc_heads.pt")) + # You'll need to handle how to load the heads in DeepSpeed case + + accelerator.register_save_state_pre_hook(save_model_hook) + accelerator.register_load_state_pre_hook(load_model_hook) + + + # Enable TF32 for faster training on Ampere GPUs, + # cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices + if args.allow_tf32 and torch.cuda.is_available(): + torch.backends.cuda.matmul.allow_tf32 = True + + if args.scale_lr: + args.learning_rate = ( + args.learning_rate * args.gradient_accumulation_steps * args.train_batch_size * accelerator.num_processes + ) + + # Use 8-bit Adam for lower memory usage or to fine-tune the model in 16GB GPUs + if args.use_8bit_adam: + try: + import bitsandbytes as bnb + except ImportError: + raise ImportError( + "To use 8-bit Adam, please install the bitsandbytes library: `pip install bitsandbytes`." + ) + + optimizer_class = bnb.optim.AdamW8bit + else: + optimizer_class = torch.optim.AdamW + + + # Optimization parameters + optimizer_G = optimizer_class( + transformer.parameters(), + lr=args.learning_rate, + betas=(args.adam_beta1, args.adam_beta2), + weight_decay=args.adam_weight_decay, + eps=args.adam_epsilon, + ) + + optimizer_D = optimizer_class( + disc.heads.parameters(), + lr=args.learning_rate, + betas=(args.adam_beta1, args.adam_beta2), + weight_decay=args.adam_weight_decay, + eps=args.adam_epsilon, + ) + + hf_dataset = load_dataset( + args.dataset_name, + data_files=args.file_path, + split='train', + ) + + train_dataset = Text2ImageDataset( + hf_dataset=hf_dataset, + resolution=args.resolution, + ) + + train_dataloader = DataLoader( + train_dataset, + batch_size=args.train_batch_size, + num_workers=args.dataloader_num_workers, + pin_memory=True, + persistent_workers=True, + shuffle=True + ) + + # Scheduler and math around the number of training steps. + overrode_max_train_steps = False + num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps) + if args.max_train_steps is None: + args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch + overrode_max_train_steps = True + + lr_scheduler = get_scheduler( + args.lr_scheduler, + optimizer=optimizer_G, + num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes, + num_training_steps=args.max_train_steps * accelerator.num_processes, + num_cycles=args.lr_num_cycles, + power=args.lr_power, + ) + + # Prepare everything with our `accelerator`. + transformer, pretrained_model, disc, optimizer_G, optimizer_D, train_dataloader, lr_scheduler = accelerator.prepare( + transformer, pretrained_model, disc, optimizer_G, optimizer_D, train_dataloader, lr_scheduler + ) + + # We need to recalculate our total training steps as the size of the training dataloader may have changed. + num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps) + if overrode_max_train_steps: + args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch + # Afterwards we recalculate our number of training epochs + args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch) + + # We need to initialize the trackers we use, and also store our configuration. + # The trackers initializes automatically on the main process. + if accelerator.is_main_process: + tracker_name = "sana-sprint" + config = { + k: str(v) if not isinstance(v, (int, float, str, bool, torch.Tensor)) else v + for k, v in vars(args).items() + } + accelerator.init_trackers(tracker_name, config=config) + + # Train! + total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps + + logger.info("***** Running training *****") + logger.info(f" Num examples = {len(train_dataset)}") + logger.info(f" Num batches each epoch = {len(train_dataloader)}") + logger.info(f" Num Epochs = {args.num_train_epochs}") + logger.info(f" Instantaneous batch size per device = {args.train_batch_size}") + logger.info(f" Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}") + logger.info(f" Gradient Accumulation steps = {args.gradient_accumulation_steps}") + logger.info(f" Total optimization steps = {args.max_train_steps}") + global_step = 0 + first_epoch = 0 + + # Potentially load in the weights and states from a previous save + if args.resume_from_checkpoint: + if args.resume_from_checkpoint != "latest": + path = os.path.basename(args.resume_from_checkpoint) + else: + # Get the mos recent checkpoint + dirs = os.listdir(args.output_dir) + dirs = [d for d in dirs if d.startswith("checkpoint")] + dirs = sorted(dirs, key=lambda x: int(x.split("-")[1])) + path = dirs[-1] if len(dirs) > 0 else None + + if path is None: + accelerator.print( + f"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run." + ) + args.resume_from_checkpoint = None + initial_global_step = 0 + else: + accelerator.print(f"Resuming from checkpoint {path}") + accelerator.load_state(os.path.join(args.output_dir, path)) + global_step = int(path.split("-")[1]) + + initial_global_step = global_step + first_epoch = global_step // num_update_steps_per_epoch + + else: + initial_global_step = 0 + + progress_bar = tqdm( + range(0, args.max_train_steps), + initial=initial_global_step, + desc="Steps", + # Only show the progress bar once on each machine. + disable=not accelerator.is_local_main_process, + ) + + phase = 'G' + vae_config_scaling_factor = vae.config.scaling_factor + sigma_data = args.sigma_data + negative_prompt = [""] * args.train_batch_size + negative_prompt_embeds, negative_prompt_attention_mask, _, _ = text_encoding_pipeline.encode_prompt( + prompt=negative_prompt, + complex_human_instruction=False, + do_classifier_free_guidance=False, + ) + + for epoch in range(first_epoch, args.num_train_epochs): + transformer.train() + disc.train() + + for step, batch in enumerate(train_dataloader): + # text encoding + prompts = batch["text"] + with torch.no_grad(): + prompt_embeds, prompt_attention_mask, _, _ = text_encoding_pipeline.encode_prompt( + prompts, complex_human_instruction=COMPLEX_HUMAN_INSTRUCTION, do_classifier_free_guidance=False + ) + + # Convert images to latent space + vae = vae.to(accelerator.device) + pixel_values = batch["image"].to(dtype=vae.dtype) + model_input = vae.encode(pixel_values).latent + model_input = model_input * vae_config_scaling_factor * sigma_data + model_input = model_input.to(dtype=weight_dtype) + + # Sample noise that we'll add to the latents + noise = torch.randn_like(model_input) * sigma_data + bsz = model_input.shape[0] + + # Sample a random timestep for each image + # for weighting schemes where we sample timesteps non-uniformly + u = compute_density_for_timestep_sampling_scm( + batch_size=bsz, + logit_mean=args.logit_mean, + logit_std=args.logit_std, + ).to(accelerator.device) + + + # Add noise according to TrigFlow. + # zt = cos(t) * x + sin(t) * noise + t = u.view(-1, 1, 1, 1) + noisy_model_input = torch.cos(t) * model_input + torch.sin(t) * noise + + + scm_cfg_scale = torch.tensor( + np.random.choice(args.scm_cfg_scale, size=bsz, replace=True), + device=accelerator.device, + ) + + def model_wrapper(scaled_x_t, t): + pred, logvar = accelerator.unwrap_model(transformer)( + hidden_states=scaled_x_t, timestep=t.flatten(), encoder_hidden_states=prompt_embeds, encoder_attention_mask=prompt_attention_mask, guidance=(scm_cfg_scale.flatten() * args.guidance_embeds_scale), jvp=True, return_logvar=True + ) + return pred, logvar + + if phase == "G": + transformer.train() + disc.eval() + models_to_accumulate = [transformer] + with accelerator.accumulate(models_to_accumulate): + with torch.no_grad(): + cfg_x_t = torch.cat([noisy_model_input, noisy_model_input], dim=0) + cfg_t = torch.cat([t, t], dim=0) + cfg_y = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) + cfg_y_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0) + + + cfg_pretrain_pred = pretrained_model( + hidden_states=(cfg_x_t / sigma_data), timestep=cfg_t.flatten(), encoder_hidden_states=cfg_y, encoder_attention_mask=cfg_y_mask + )[0] + + cfg_dxt_dt = sigma_data * cfg_pretrain_pred + + dxt_dt_uncond, dxt_dt = cfg_dxt_dt.chunk(2) + + scm_cfg_scale = scm_cfg_scale.view(-1, 1, 1, 1) + dxt_dt = dxt_dt_uncond + scm_cfg_scale * (dxt_dt - dxt_dt_uncond) + + + v_x = torch.cos(t) * torch.sin(t) * dxt_dt / sigma_data + v_t = torch.cos(t) * torch.sin(t) + + + # Adapt from https://github.com/xandergos/sCM-mnist/blob/master/train_consistency.py + with torch.no_grad(): + F_theta, F_theta_grad, logvar = torch.func.jvp( + model_wrapper, (noisy_model_input / sigma_data, t), (v_x, v_t), has_aux=True + ) + + + F_theta, logvar = transformer( + hidden_states = (noisy_model_input / sigma_data), + timestep=t.flatten(), + encoder_hidden_states=prompt_embeds, + encoder_attention_mask=prompt_attention_mask, + guidance=(scm_cfg_scale.flatten() * args.guidance_embeds_scale), + return_logvar=True + ) + + + logvar = logvar.view(-1, 1, 1, 1) + F_theta_grad = F_theta_grad.detach() + F_theta_minus = F_theta.detach() + + # Warmup steps + r = min(1, global_step / args.tangent_warmup_steps) + + # Calculate gradient g using JVP rearrangement + g = -torch.cos(t) * torch.cos(t) * (sigma_data * F_theta_minus - dxt_dt) + second_term = -r * (torch.cos(t) * torch.sin(t) * noisy_model_input + sigma_data * F_theta_grad) + g = g + second_term + + # Tangent normalization + g_norm = torch.linalg.vector_norm(g, dim=(1, 2, 3), keepdim=True) + g = g / (g_norm + 0.1) # 0.1 is the constant c, can be modified but 0.1 was used in the paper + + sigma = torch.tan(t) * sigma_data + weight = 1 / sigma + + l2_loss = torch.square(F_theta - F_theta_minus - g) + + # Calculate loss with normalization factor + loss = (weight / torch.exp(logvar)) * l2_loss + logvar + + loss = loss.mean() + + loss_no_logvar = weight * torch.square(F_theta - F_theta_minus - g) + loss_no_logvar = loss_no_logvar.mean() + loss_no_weight = l2_loss.mean() + g_norm = g_norm.mean() + + + pred_x_0 = torch.cos(t) * noisy_model_input - torch.sin(t) * F_theta * sigma_data + + if args.train_largest_timestep: + pred_x_0.detach() + u = compute_density_for_timestep_sampling_scm( + batch_size=bsz, + logit_mean=args.logit_mean, + logit_std=args.logit_std, + ).to(accelerator.device) + t_new = u.view(-1, 1, 1, 1) + + random_mask = torch.rand_like(t_new) < args.largest_timestep_prob + + t_new = torch.where(random_mask, torch.full_like(t_new, args.largest_timestep), t_new) + z_new = torch.randn_like(model_input) * sigma_data + x_t_new = torch.cos(t_new) * model_input + torch.sin(t_new) * z_new + + F_theta = transformer( + hidden_states = (x_t_new / sigma_data), + timestep=t_new.flatten(), + encoder_hidden_states=prompt_embeds, + encoder_attention_mask=prompt_attention_mask, + guidance=(scm_cfg_scale.flatten() * args.guidance_embeds_scale), + return_logvar=False, + jvp=False, + )[0] + + + pred_x_0 = torch.cos(t_new) * x_t_new - torch.sin(t_new) * F_theta * sigma_data + + # Sample timesteps for discriminator + timesteps_D = compute_density_for_timestep_sampling_scm( + batch_size=bsz, + logit_mean=args.logit_mean_discriminator, + logit_std=args.logit_std_discriminator, + ).to(accelerator.device) + t_D = timesteps_D.view(-1, 1, 1, 1) + + # Add noise to predicted x0 + z_D = torch.randn_like(model_input) * sigma_data + noised_predicted_x0 = torch.cos(t_D) * pred_x_0 + torch.sin(t_D) * z_D + + + # Calculate adversarial loss + pred_fake = disc(hidden_states=(noised_predicted_x0 / sigma_data), timestep=t_D.flatten(), encoder_hidden_states=prompt_embeds, encoder_attention_mask=prompt_attention_mask) + adv_loss = -torch.mean(pred_fake) + + + # Total loss = sCM loss + LADD loss + + total_loss = args.scm_lambda * loss + adv_loss * args.adv_lambda + + total_loss = total_loss / args.gradient_accumulation_steps + + accelerator.backward(total_loss) + + + + if accelerator.sync_gradients: + grad_norm = accelerator.clip_grad_norm_(transformer.parameters(), args.gradient_clip) + if torch.logical_or(grad_norm.isnan(), grad_norm.isinf()): + optimizer_G.zero_grad(set_to_none=True) + optimizer_D.zero_grad(set_to_none=True) + logger.warning("NaN or Inf detected in grad_norm, skipping iteration...") + continue + + # switch phase to D + phase = "D" + + optimizer_G.step() + lr_scheduler.step() + optimizer_G.zero_grad(set_to_none=True) + + elif phase == "D": + transformer.eval() + disc.train() + models_to_accumulate = [disc] + with accelerator.accumulate(models_to_accumulate): + with torch.no_grad(): + scm_cfg_scale = torch.tensor( + np.random.choice(args.scm_cfg_scale, size=bsz, replace=True), + device=accelerator.device, + ) + + if args.train_largest_timestep: + random_mask = torch.rand_like(t) < args.largest_timestep_prob + t = torch.where(random_mask, torch.full_like(t, args.largest_timestep_prob), t) + + z_new = torch.randn_like(model_input) * sigma_data + noisy_model_input = torch.cos(t) * model_input + torch.sin(t) * z_new + # here + F_theta = transformer( + hidden_states = (noisy_model_input / sigma_data), + timestep=t.flatten(), + encoder_hidden_states=prompt_embeds, + encoder_attention_mask=prompt_attention_mask, + guidance=(scm_cfg_scale.flatten() * args.guidance_embeds_scale), + return_logvar=False, + jvp=False, + )[0] + pred_x_0 = torch.cos(t) * noisy_model_input - torch.sin(t) * F_theta * sigma_data + + # Sample timesteps for fake and real samples + timestep_D_fake = compute_density_for_timestep_sampling_scm( + batch_size=bsz, + logit_mean=args.logit_mean_discriminator, + logit_std=args.logit_std_discriminator, + ).to(accelerator.device) + timesteps_D_real = timestep_D_fake + + t_D_fake = timestep_D_fake.view(-1, 1, 1, 1) + t_D_real = timesteps_D_real.view(-1, 1, 1, 1) + + # Add noise to predicted x0 and real x0 + z_D_fake = torch.randn_like(model_input) * sigma_data + z_D_real = torch.randn_like(model_input) * sigma_data + noised_predicted_x0 = torch.cos(t_D_fake) * pred_x_0 + torch.sin(t_D_fake) * z_D_fake + noised_latents = torch.cos(t_D_real) * model_input + torch.sin(t_D_real) * z_D_real + + # Add misaligned pairs if enabled and batch size > 1 + if args.misaligned_pairs_D and bsz > 1: + # Create shifted pairs + shifted_x0 = torch.roll(model_input, 1, 0) + timesteps_D_shifted = compute_density_for_timestep_sampling_scm( + batch_size=bsz, + logit_mean=args.logit_mean_discriminator, + logit_std=args.logit_std_discriminator, + ).to(accelerator.device) + t_D_shifted = timesteps_D_shifted.view(-1, 1, 1, 1) + + # Add noise to shifted pairs + z_D_shifted = torch.randn_like(shifted_x0) * sigma_data + noised_shifted_x0 = torch.cos(t_D_shifted) * shifted_x0 + torch.sin(t_D_shifted) * z_D_shifted + + # Concatenate with original noised samples + noised_predicted_x0 = torch.cat([noised_predicted_x0, noised_shifted_x0], dim=0) + t_D_fake = torch.cat([t_D_fake, t_D_shifted], dim=0) + prompt_embeds = torch.cat([prompt_embeds, prompt_embeds], dim=0) + prompt_attention_mask = torch.cat([prompt_attention_mask, prompt_attention_mask], dim=0) + + + # Calculate D loss + + pred_fake = disc(hidden_states=(noised_predicted_x0 / sigma_data), timestep=t_D_fake.flatten(), encoder_hidden_states=prompt_embeds, encoder_attention_mask=prompt_attention_mask) + pred_true = disc(hidden_states=(noised_latents / sigma_data), timestep=t_D_real.flatten(), encoder_hidden_states=prompt_embeds, encoder_attention_mask=prompt_attention_mask) + + # hinge loss + loss_real = torch.mean(F.relu(1.0 - pred_true)) + loss_gen = torch.mean(F.relu(1.0 + pred_fake)) + loss_D = 0.5 * (loss_real + loss_gen) + + loss_D = loss_D / args.gradient_accumulation_steps + + accelerator.backward(loss_D) + + + if accelerator.sync_gradients: + grad_norm = accelerator.clip_grad_norm_(disc.parameters(), args.gradient_clip) + if torch.logical_or(grad_norm.isnan(), grad_norm.isinf()): + optimizer_G.zero_grad(set_to_none=True) + optimizer_D.zero_grad(set_to_none=True) + logger.warning("NaN or Inf detected in grad_norm, skipping iteration...") + continue + + # switch back to phase G and add global step by one. + phase = "G" + + optimizer_D.step() + optimizer_D.zero_grad(set_to_none=True) + + + # Checks if the accelerator has performed an optimization step behind the scenes + if accelerator.sync_gradients: + progress_bar.update(1) + global_step += 1 + + if accelerator.is_main_process: + if global_step % args.checkpointing_steps == 0: + # _before_ saving state, check if this save would set us over the `checkpoints_total_limit` + if args.checkpoints_total_limit is not None: + checkpoints = os.listdir(args.output_dir) + checkpoints = [d for d in checkpoints if d.startswith("checkpoint")] + checkpoints = sorted(checkpoints, key=lambda x: int(x.split("-")[1])) + + # before we save the new checkpoint, we need to have at _most_ `checkpoints_total_limit - 1` checkpoints + if len(checkpoints) >= args.checkpoints_total_limit: + num_to_remove = len(checkpoints) - args.checkpoints_total_limit + 1 + removing_checkpoints = checkpoints[0:num_to_remove] + + logger.info( + f"{len(checkpoints)} checkpoints already exist, removing {len(removing_checkpoints)} checkpoints" + ) + logger.info(f"removing checkpoints: {', '.join(removing_checkpoints)}") + + for removing_checkpoint in removing_checkpoints: + removing_checkpoint = os.path.join(args.output_dir, removing_checkpoint) + shutil.rmtree(removing_checkpoint) + + save_path = os.path.join(args.output_dir, f"checkpoint-{global_step}") + accelerator.save_state(save_path) + logger.info(f"Saved state to {save_path}") + + logs = {"scm_loss": loss.detach().item(), "adv_loss": adv_loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]} + progress_bar.set_postfix(**logs) + accelerator.log(logs, step=global_step) + + if global_step >= args.max_train_steps: + break + + if accelerator.is_main_process: + if args.validation_prompt is not None and epoch % args.validation_epochs == 0: + # create pipeline + pipeline = SanaSprintPipeline.from_pretrained( + args.pretrained_model_name_or_path, + transformer=accelerator.unwrap_model(transformer), + revision=args.revision, + variant=args.variant, + torch_dtype=torch.float32, + ) + pipeline_args = { + "prompt": args.validation_prompt, + "complex_human_instruction": COMPLEX_HUMAN_INSTRUCTION, + } + images = log_validation( + pipeline=pipeline, + args=args, + accelerator=accelerator, + pipeline_args=pipeline_args, + epoch=epoch, + ) + free_memory() + + images = None + del pipeline + + + accelerator.wait_for_everyone() + if accelerator.is_main_process: + transformer = unwrap_model(transformer) + if args.upcast_before_saving: + transformer.to(torch.float32) + else: + transformer = transformer.to(weight_dtype) + + # Save discriminator heads + disc = unwrap_model(disc) + disc_heads_state_dict = disc.heads.state_dict() + + # Save transformer model + transformer.save_pretrained(os.path.join(args.output_dir, "transformer")) + + # Save discriminator heads + torch.save(disc_heads_state_dict, os.path.join(args.output_dir, "disc_heads.pt")) + + # Final inference + # Load previous pipeline + pipeline = SanaSprintPipeline.from_pretrained( + args.pretrained_model_name_or_path, + transformer=accelerator.unwrap_model(transformer), + revision=args.revision, + variant=args.variant, + torch_dtype=torch.float32, + ) + + # run inference + images = [] + if args.validation_prompt and args.num_validation_images > 0: + pipeline_args = { + "prompt": args.validation_prompt, + "complex_human_instruction": COMPLEX_HUMAN_INSTRUCTION, + } + images = log_validation( + pipeline=pipeline, + args=args, + accelerator=accelerator, + pipeline_args=pipeline_args, + epoch=epoch, + is_final_validation=True, + ) + + if args.push_to_hub: + save_model_card( + repo_id, + images=images, + base_model=args.pretrained_model_name_or_path, + instance_prompt=args.instance_prompt, + validation_prompt=args.validation_prompt, + repo_folder=args.output_dir, + ) + upload_folder( + repo_id=repo_id, + folder_path=args.output_dir, + commit_message="End of training", + ignore_patterns=["step_*", "epoch_*"], + ) + + images = None + del pipeline + + accelerator.end_training() + + +if __name__ == "__main__": + args = parse_args() + main(args) \ No newline at end of file diff --git a/examples/research_projects/sana/train_sana_sprint_diffusers.sh b/examples/research_projects/sana/train_sana_sprint_diffusers.sh new file mode 100644 index 000000000000..bf6f47f9d731 --- /dev/null +++ b/examples/research_projects/sana/train_sana_sprint_diffusers.sh @@ -0,0 +1,24 @@ +your_local_path='output' + +huggingface-cli download Efficient-Large-Model/SANA_Sprint_1.6B_1024px_teacher_diffusers --local-dir $your_local_path/SANA_Sprint_1.6B_1024px_teacher_diffusers + +python train_sana_sprint_diffusers.py \ + --pretrained_model_name_or_path=$your_local_path/SANA_Sprint_1.6B_1024px_teacher_diffusers \ + --output_dir=$your_local_path \ + --mixed_precision=bf16 \ + --resolution=1024 \ + --learning_rate=1e-6 \ + --max_train_steps=30000 \ + --dataloader_num_workers=8 \ + --dataset_name='brivangl/midjourney-v6-llava' \ + --file_path data/train_000.parquet data/train_001.parquet data/train_002.parquet \ + --checkpointing_steps=500 --checkpoints_total_limit=10 \ + --train_batch_size=1 \ + --gradient_accumulation_steps=1 \ + --seed=453645634 \ + --train_largest_timestep \ + --misaligned_pairs_D \ + --gradient_checkpointing \ + --resume_from_checkpoint="latest" \ + + From 9cb050b2efc261bc9231a0a17b4d67c36caacc48 Mon Sep 17 00:00:00 2001 From: lawrence-cj Date: Wed, 7 May 2025 17:50:14 +0800 Subject: [PATCH 04/10] make style && make quality; --- .../sana/train_sana_sprint_diffusers.py | 98 +++++++------------ .../models/transformers/sana_transformer.py | 4 +- 2 files changed, 39 insertions(+), 63 deletions(-) diff --git a/examples/research_projects/sana/train_sana_sprint_diffusers.py b/examples/research_projects/sana/train_sana_sprint_diffusers.py index 15aa323953b1..1efac298f143 100644 --- a/examples/research_projects/sana/train_sana_sprint_diffusers.py +++ b/examples/research_projects/sana/train_sana_sprint_diffusers.py @@ -15,72 +15,48 @@ # See the License for the specific language governing permissions and import argparse -import copy -import itertools -import json +import io import logging import math import os -import random import shutil -import warnings from pathlib import Path +from typing import Callable import accelerate -import io import numpy as np import torch import torch.nn as nn import torch.nn.functional as F import torch.utils.checkpoint import torchvision.transforms as T -import torchvision.transforms.functional as TF import transformers -import webdataset as wds from accelerate import Accelerator from accelerate.logging import get_logger from accelerate.utils import DistributedDataParallelKwargs, DistributedType, ProjectConfiguration, set_seed -from braceexpand import braceexpand from datasets import load_dataset from huggingface_hub import create_repo, upload_folder -from huggingface_hub.utils import insecure_hashlib from packaging import version -from peft.utils import get_peft_model_state_dict from PIL import Image -from PIL.ImageOps import exif_transpose from safetensors.torch import load_file from torch.nn.utils.spectral_norm import SpectralNorm -from torch.utils.data import default_collate, Dataset, DataLoader -from torchvision.transforms.functional import crop +from torch.utils.data import DataLoader, Dataset from tqdm.auto import tqdm from transformers import AutoTokenizer, Gemma2Model -from typing import Callable, List, Union -from webdataset.tariterators import ( - base_plus_ext, - tar_file_expander, - url_opener, - valid_sample, -) import diffusers from diffusers import ( AutoencoderDC, - FlowMatchEulerDiscreteScheduler, SanaPipeline, SanaSprintPipeline, SanaTransformer2DModel, - SCMScheduler, ) from diffusers.optimization import get_scheduler from diffusers.training_utils import ( - cast_training_params, - compute_density_for_timestep_sampling, - compute_loss_weighting_for_sd3, free_memory, ) from diffusers.utils import ( check_min_version, - convert_unet_state_dict_to_peft, is_wandb_available, ) from diffusers.utils.hub_utils import load_or_create_model_card, populate_model_card @@ -98,7 +74,7 @@ if is_torch_npu_available(): torch.npu.config.allow_internal_format = False - + COMPLEX_HUMAN_INSTRUCTION = [ "Given a user prompt, generate an 'Enhanced prompt' that provides detailed visual descriptions suitable for image generation. Evaluate the level of detail in the user prompt:", "- If the prompt is simple, focus on adding specifics about colors, shapes, sizes, textures, and spatial relationships to create vivid and concrete scenes.", @@ -109,7 +85,7 @@ "Please generate only the enhanced description for the prompt below and avoid including any additional commentary or evaluations:", "User Prompt: ", ] - + class Text2ImageDataset(Dataset): @@ -140,17 +116,17 @@ def __init__(self, hf_dataset, resolution=1024): def __len__(self): return len(self.dataset) - + def __getitem__(self, idx): item = self.dataset[idx] text = item['llava'] image_bytes = item['image'] - + # Convert bytes to PIL Image image = Image.open(io.BytesIO(image_bytes)) - + image_tensor = self.transform(image) - + return { 'text': text, 'image': image_tensor @@ -768,7 +744,7 @@ def state_dict(self): def __getattr__(self, name): return getattr(self.disc, name) - + class SanaTrigFlow(SanaTransformer2DModel): def __init__(self, original_model, guidance=False): self.__dict__ = original_model.__dict__ @@ -779,7 +755,7 @@ def __init__(self, original_model, guidance=False): self.logvar_linear = torch.nn.Linear(hidden_size, 1) torch.nn.init.xavier_uniform_(self.logvar_linear.weight) torch.nn.init.constant_(self.logvar_linear.bias, 0) - + def forward(self, hidden_states, encoder_hidden_states, timestep, guidance=None, jvp=False, return_logvar=False, **kwargs): batch_size = hidden_states.shape[0] latents = hidden_states @@ -812,8 +788,8 @@ def forward(self, hidden_states, encoder_hidden_states, timestep, guidance=None, trigflow_model_out = ((1 - 2 * flow_timestep_expanded) * latent_model_input + (1 - 2 * flow_timestep_expanded + 2 * flow_timestep_expanded**2) * model_out) / torch.sqrt( flow_timestep_expanded**2 + (1 - flow_timestep_expanded) ** 2 ) - - + + if self.guidance and guidance is not None: timestep, embedded_timestep = self.time_embed( timestep, guidance=guidance, hidden_dtype=hidden_states.dtype @@ -822,15 +798,15 @@ def forward(self, hidden_states, encoder_hidden_states, timestep, guidance=None, timestep, embedded_timestep = self.time_embed( timestep, batch_size=batch_size, hidden_dtype=hidden_states.dtype ) - + if return_logvar: logvar = self.logvar_linear(embedded_timestep) return trigflow_model_out, logvar - + return (trigflow_model_out,) - + def compute_density_for_timestep_sampling_scm( batch_size: int, logit_mean: float = None, logit_std: float = None @@ -925,19 +901,19 @@ def main(args): revision=args.revision, variant=args.variant, ) - + ori_transformer = SanaTransformer2DModel.from_pretrained( args.pretrained_model_name_or_path, subfolder="transformer", revision=args.revision, variant=args.variant, guidance_embeds=True, cross_attention_type='vanilla' ) - + ori_transformer_no_guide = SanaTransformer2DModel.from_pretrained( args.pretrained_model_name_or_path, subfolder="transformer", revision=args.revision, variant=args.variant, guidance_embeds=False ) - + original_state_dict = load_file(f"{args.pretrained_model_name_or_path}/transformer/diffusion_pytorch_model.safetensors") - + param_mapping = { 'time_embed.emb.timestep_embedder.linear_1.weight': 'time_embed.timestep_embedder.linear_1.weight', 'time_embed.emb.timestep_embedder.linear_1.bias': 'time_embed.timestep_embedder.linear_1.bias', @@ -968,7 +944,7 @@ def main(args): transformer = SanaTrigFlow(ori_transformer, guidance=True).train() pretrained_model = SanaTrigFlow(ori_transformer_no_guide, guidance=False).eval() - + disc = SanaMSCMDiscriminator( pretrained_model, is_multiscale=args.ladd_multi_scale, @@ -1134,7 +1110,7 @@ def load_model_hook(models, input_dir): data_files=args.file_path, split='train', ) - + train_dataset = Text2ImageDataset( hf_dataset=hf_dataset, resolution=args.resolution, @@ -1282,8 +1258,8 @@ def load_model_hook(models, input_dir): # Add noise according to TrigFlow. # zt = cos(t) * x + sin(t) * noise t = u.view(-1, 1, 1, 1) - noisy_model_input = torch.cos(t) * model_input + torch.sin(t) * noise - + noisy_model_input = torch.cos(t) * model_input + torch.sin(t) * noise + scm_cfg_scale = torch.tensor( np.random.choice(args.scm_cfg_scale, size=bsz, replace=True), @@ -1295,7 +1271,7 @@ def model_wrapper(scaled_x_t, t): hidden_states=scaled_x_t, timestep=t.flatten(), encoder_hidden_states=prompt_embeds, encoder_attention_mask=prompt_attention_mask, guidance=(scm_cfg_scale.flatten() * args.guidance_embeds_scale), jvp=True, return_logvar=True ) return pred, logvar - + if phase == "G": transformer.train() disc.eval() @@ -1322,8 +1298,8 @@ def model_wrapper(scaled_x_t, t): v_x = torch.cos(t) * torch.sin(t) * dxt_dt / sigma_data v_t = torch.cos(t) * torch.sin(t) - - + + # Adapt from https://github.com/xandergos/sCM-mnist/blob/master/train_consistency.py with torch.no_grad(): F_theta, F_theta_grad, logvar = torch.func.jvp( @@ -1371,8 +1347,8 @@ def model_wrapper(scaled_x_t, t): loss_no_logvar = loss_no_logvar.mean() loss_no_weight = l2_loss.mean() g_norm = g_norm.mean() - - + + pred_x_0 = torch.cos(t) * noisy_model_input - torch.sin(t) * F_theta * sigma_data if args.train_largest_timestep: @@ -1414,7 +1390,7 @@ def model_wrapper(scaled_x_t, t): # Add noise to predicted x0 z_D = torch.randn_like(model_input) * sigma_data noised_predicted_x0 = torch.cos(t_D) * pred_x_0 + torch.sin(t_D) * z_D - + # Calculate adversarial loss pred_fake = disc(hidden_states=(noised_predicted_x0 / sigma_data), timestep=t_D.flatten(), encoder_hidden_states=prompt_embeds, encoder_attention_mask=prompt_attention_mask) @@ -1445,7 +1421,7 @@ def model_wrapper(scaled_x_t, t): optimizer_G.step() lr_scheduler.step() optimizer_G.zero_grad(set_to_none=True) - + elif phase == "D": transformer.eval() disc.train() @@ -1515,7 +1491,7 @@ def model_wrapper(scaled_x_t, t): # Calculate D loss - + pred_fake = disc(hidden_states=(noised_predicted_x0 / sigma_data), timestep=t_D_fake.flatten(), encoder_hidden_states=prompt_embeds, encoder_attention_mask=prompt_attention_mask) pred_true = disc(hidden_states=(noised_latents / sigma_data), timestep=t_D_real.flatten(), encoder_hidden_states=prompt_embeds, encoder_attention_mask=prompt_attention_mask) @@ -1542,7 +1518,7 @@ def model_wrapper(scaled_x_t, t): optimizer_D.step() optimizer_D.zero_grad(set_to_none=True) - + # Checks if the accelerator has performed an optimization step behind the scenes if accelerator.sync_gradients: @@ -1616,14 +1592,14 @@ def model_wrapper(scaled_x_t, t): transformer.to(torch.float32) else: transformer = transformer.to(weight_dtype) - + # Save discriminator heads disc = unwrap_model(disc) disc_heads_state_dict = disc.heads.state_dict() - + # Save transformer model transformer.save_pretrained(os.path.join(args.output_dir, "transformer")) - + # Save discriminator heads torch.save(disc_heads_state_dict, os.path.join(args.output_dir, "disc_heads.pt")) @@ -1677,4 +1653,4 @@ def model_wrapper(scaled_x_t, t): if __name__ == "__main__": args = parse_args() - main(args) \ No newline at end of file + main(args) diff --git a/src/diffusers/models/transformers/sana_transformer.py b/src/diffusers/models/transformers/sana_transformer.py index 1d83283882a4..948bb7eeba07 100644 --- a/src/diffusers/models/transformers/sana_transformer.py +++ b/src/diffusers/models/transformers/sana_transformer.py @@ -12,9 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. +import math from typing import Any, Dict, Optional, Tuple, Union -import math import torch import torch.nn.functional as F from torch import nn @@ -185,7 +185,7 @@ def __call__( return hidden_states - + class SanaAttnProcessor3_0: r""" Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0). From 86bef58618c5f929c44ec55f7f369cd6f116efdb Mon Sep 17 00:00:00 2001 From: lawrence-cj Date: Thu, 8 May 2025 15:25:27 +0800 Subject: [PATCH 05/10] modify the attention processor with `set_attn_processor` and change `SanaAttnProcessor3_0` to `SanaVanillaAttnProcessor` --- .../sana/train_sana_sprint_diffusers.py | 102 ++++++++++++++---- .../models/transformers/sana_transformer.py | 97 +---------------- 2 files changed, 84 insertions(+), 115 deletions(-) diff --git a/examples/research_projects/sana/train_sana_sprint_diffusers.py b/examples/research_projects/sana/train_sana_sprint_diffusers.py index 1efac298f143..2503e78636bf 100644 --- a/examples/research_projects/sana/train_sana_sprint_diffusers.py +++ b/examples/research_projects/sana/train_sana_sprint_diffusers.py @@ -86,6 +86,86 @@ "User Prompt: ", ] +class SanaVanillaAttnProcessor: + r""" + Processor for implementing scaled dot-product attention to support JVP calculation during training. + """ + + def __init__(self): + pass + + @staticmethod + def scaled_dot_product_attention(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, scale=None + ) -> torch.Tensor: + B, H, L, S = *query.size()[:-1], key.size(-2) + scale_factor = 1 / math.sqrt(query.size(-1)) if scale is None else scale + attn_bias = torch.zeros(B, H, L, S, dtype=query.dtype, device=query.device) + + if attn_mask is not None: + if attn_mask.dtype == torch.bool: + attn_bias.masked_fill_(attn_mask.logical_not(), float("-inf")) + else: + attn_bias += attn_mask + attn_weight = query @ key.transpose(-2, -1) * scale_factor + attn_weight += attn_bias + attn_weight = torch.softmax(attn_weight, dim=-1) + attn_weight = torch.dropout(attn_weight, dropout_p, train=True) + return attn_weight @ value + + def __call__( + self, + attn: Attention, + hidden_states: torch.Tensor, + encoder_hidden_states: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + batch_size, sequence_length, _ = ( + hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape + ) + + if attention_mask is not None: + attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size) + # scaled_dot_product_attention expects attention_mask shape to be + # (batch, heads, source_length, target_length) + attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1]) + + query = attn.to_q(hidden_states) + + if encoder_hidden_states is None: + encoder_hidden_states = hidden_states + + key = attn.to_k(encoder_hidden_states) + value = attn.to_v(encoder_hidden_states) + + if attn.norm_q is not None: + query = attn.norm_q(query) + if attn.norm_k is not None: + key = attn.norm_k(key) + + inner_dim = key.shape[-1] + head_dim = inner_dim // attn.heads + + query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + + key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + + # the output of sdp = (batch, num_heads, seq_len, head_dim) + hidden_states = self.scaled_dot_product_attention( + query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False + ) + + hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) + hidden_states = hidden_states.to(query.dtype) + + # linear proj + hidden_states = attn.to_out[0](hidden_states) + # dropout + hidden_states = attn.to_out[1](hidden_states) + + hidden_states = hidden_states / attn.rescale_output_factor + + return hidden_states class Text2ImageDataset(Dataset): @@ -109,7 +189,6 @@ def __init__(self, hf_dataset, resolution=1024): T.Lambda(lambda img: img.convert("RGB")), T.Resize(resolution), # Image.BICUBIC T.CenterCrop(resolution), - # T.RandomHorizontalFlip(), T.ToTensor(), T.Normalize([0.5], [0.5]), ]) @@ -132,7 +211,7 @@ def __getitem__(self, idx): 'image': image_tensor } -# TODO here + def save_model_card( repo_id: str, images=None, @@ -807,7 +886,6 @@ def forward(self, hidden_states, encoder_hidden_states, timestep, guidance=None, return (trigflow_model_out,) - def compute_density_for_timestep_sampling_scm( batch_size: int, logit_mean: float = None, logit_std: float = None ): @@ -820,7 +898,6 @@ def compute_density_for_timestep_sampling_scm( return u - def main(args): if args.report_to == "wandb" and args.hub_token is not None: raise ValueError( @@ -872,7 +949,6 @@ def main(args): if args.seed is not None: set_seed(args.seed) - # Handle the repository creation if accelerator.is_main_process: if args.output_dir is not None: @@ -904,8 +980,9 @@ def main(args): ori_transformer = SanaTransformer2DModel.from_pretrained( args.pretrained_model_name_or_path, subfolder="transformer", revision=args.revision, variant=args.variant, - guidance_embeds=True, cross_attention_type='vanilla' + guidance_embeds=True, ) + ori_transformer.set_attn_processor(SanaVanillaAttnProcessor()) ori_transformer_no_guide = SanaTransformer2DModel.from_pretrained( args.pretrained_model_name_or_path, subfolder="transformer", revision=args.revision, variant=args.variant, @@ -929,7 +1006,6 @@ def main(args): zero_state_dict = {} - target_device = accelerator.device param_w1 = guidance_embedder_module.linear_1.weight zero_state_dict['linear_1.weight'] = torch.zeros(param_w1.shape, device=target_device) @@ -941,7 +1017,6 @@ def main(args): zero_state_dict['linear_2.bias'] = torch.zeros(param_b2.shape, device=target_device) guidance_embedder_module.load_state_dict(zero_state_dict, strict=False, assign=True) - transformer = SanaTrigFlow(ori_transformer, guidance=True).train() pretrained_model = SanaTrigFlow(ori_transformer_no_guide, guidance=False).eval() @@ -951,7 +1026,6 @@ def main(args): head_block_ids=args.head_block_ids, ).train() - transformer.requires_grad_(True) pretrained_model.requires_grad_(False) disc.model.requires_grad_(False) @@ -1005,7 +1079,6 @@ def main(args): if args.gradient_checkpointing: transformer.enable_gradient_checkpointing() - def unwrap_model(model): model = accelerator.unwrap_model(model) model = model._orig_mod if is_compiled_module(model) else model @@ -1063,7 +1136,6 @@ def load_model_hook(models, input_dir): accelerator.register_save_state_pre_hook(save_model_hook) accelerator.register_load_state_pre_hook(load_model_hook) - # Enable TF32 for faster training on Ampere GPUs, # cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices if args.allow_tf32 and torch.cuda.is_available(): @@ -1087,7 +1159,6 @@ def load_model_hook(models, input_dir): else: optimizer_class = torch.optim.AdamW - # Optimization parameters optimizer_G = optimizer_class( transformer.parameters(), @@ -1391,12 +1462,10 @@ def model_wrapper(scaled_x_t, t): z_D = torch.randn_like(model_input) * sigma_data noised_predicted_x0 = torch.cos(t_D) * pred_x_0 + torch.sin(t_D) * z_D - # Calculate adversarial loss pred_fake = disc(hidden_states=(noised_predicted_x0 / sigma_data), timestep=t_D.flatten(), encoder_hidden_states=prompt_embeds, encoder_attention_mask=prompt_attention_mask) adv_loss = -torch.mean(pred_fake) - # Total loss = sCM loss + LADD loss total_loss = args.scm_lambda * loss + adv_loss * args.adv_lambda @@ -1405,8 +1474,6 @@ def model_wrapper(scaled_x_t, t): accelerator.backward(total_loss) - - if accelerator.sync_gradients: grad_norm = accelerator.clip_grad_norm_(transformer.parameters(), args.gradient_clip) if torch.logical_or(grad_norm.isnan(), grad_norm.isinf()): @@ -1504,7 +1571,6 @@ def model_wrapper(scaled_x_t, t): accelerator.backward(loss_D) - if accelerator.sync_gradients: grad_norm = accelerator.clip_grad_norm_(disc.parameters(), args.gradient_clip) if torch.logical_or(grad_norm.isnan(), grad_norm.isinf()): @@ -1519,7 +1585,6 @@ def model_wrapper(scaled_x_t, t): optimizer_D.step() optimizer_D.zero_grad(set_to_none=True) - # Checks if the accelerator has performed an optimization step behind the scenes if accelerator.sync_gradients: progress_bar.update(1) @@ -1584,7 +1649,6 @@ def model_wrapper(scaled_x_t, t): images = None del pipeline - accelerator.wait_for_everyone() if accelerator.is_main_process: transformer = unwrap_model(transformer) diff --git a/src/diffusers/models/transformers/sana_transformer.py b/src/diffusers/models/transformers/sana_transformer.py index 948bb7eeba07..52236275dc88 100644 --- a/src/diffusers/models/transformers/sana_transformer.py +++ b/src/diffusers/models/transformers/sana_transformer.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import math from typing import Any, Dict, Optional, Tuple, Union import torch @@ -186,91 +185,6 @@ def __call__( return hidden_states -class SanaAttnProcessor3_0: - r""" - Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0). - """ - - def __init__(self): - if not hasattr(F, "scaled_dot_product_attention"): - raise ImportError("SanaAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.") - - @staticmethod - def scaled_dot_product_attention(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, scale=None - ) -> torch.Tensor: - B, H, L, S = *query.size()[:-1], key.size(-2) - scale_factor = 1 / math.sqrt(query.size(-1)) if scale is None else scale - attn_bias = torch.zeros(B, H, L, S, dtype=query.dtype, device=query.device) - - if attn_mask is not None: - if attn_mask.dtype == torch.bool: - attn_bias.masked_fill_(attn_mask.logical_not(), float("-inf")) - else: - attn_bias += attn_mask - attn_weight = query @ key.transpose(-2, -1) * scale_factor - attn_weight += attn_bias - attn_weight = torch.softmax(attn_weight, dim=-1) - attn_weight = torch.dropout(attn_weight, dropout_p, train=True) - return attn_weight @ value - - # return x - def __call__( - self, - attn: Attention, - hidden_states: torch.Tensor, - encoder_hidden_states: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - batch_size, sequence_length, _ = ( - hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape - ) - - if attention_mask is not None: - attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size) - # scaled_dot_product_attention expects attention_mask shape to be - # (batch, heads, source_length, target_length) - attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1]) - - query = attn.to_q(hidden_states) - - if encoder_hidden_states is None: - encoder_hidden_states = hidden_states - - key = attn.to_k(encoder_hidden_states) - value = attn.to_v(encoder_hidden_states) - - if attn.norm_q is not None: - query = attn.norm_q(query) - if attn.norm_k is not None: - key = attn.norm_k(key) - - inner_dim = key.shape[-1] - head_dim = inner_dim // attn.heads - - query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - - key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - - # the output of sdp = (batch, num_heads, seq_len, head_dim) - # TODO: add support for attn.scale when we move to Torch 2.1 - hidden_states = self.scaled_dot_product_attention( - query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False - ) - - hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) - hidden_states = hidden_states.to(query.dtype) - - # linear proj - hidden_states = attn.to_out[0](hidden_states) - # dropout - hidden_states = attn.to_out[1](hidden_states) - - hidden_states = hidden_states / attn.rescale_output_factor - - return hidden_states - - class SanaTransformerBlock(nn.Module): r""" Transformer block introduced in [Sana](https://huggingface.co/papers/2410.10629). @@ -291,7 +205,6 @@ def __init__( attention_out_bias: bool = True, mlp_ratio: float = 2.5, qk_norm: Optional[str] = None, - cross_attention_type: str = "flash", ) -> None: super().__init__() @@ -310,12 +223,6 @@ def __init__( ) # 2. Cross Attention - if cross_attention_type == "flash": - cross_attention_processor = SanaAttnProcessor2_0() - elif cross_attention_type == "vanilla": - cross_attention_processor = SanaAttnProcessor3_0() - else: - raise ValueError(f"Cross attention type {cross_attention_type} is not defined.") if cross_attention_dim is not None: self.norm2 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps) self.attn2 = Attention( @@ -328,7 +235,7 @@ def __init__( dropout=dropout, bias=True, out_bias=attention_out_bias, - processor=cross_attention_processor, + processor=SanaAttnProcessor2_0(), ) # 3. Feed-forward @@ -453,7 +360,6 @@ def __init__( guidance_embeds_scale: float = 0.1, qk_norm: Optional[str] = None, timestep_scale: float = 1.0, - cross_attention_type: str = "flash", ) -> None: super().__init__() @@ -496,7 +402,6 @@ def __init__( norm_eps=norm_eps, mlp_ratio=mlp_ratio, qk_norm=qk_norm, - cross_attention_type=cross_attention_type, ) for _ in range(num_layers) ] From 6c3a398a37b9bc7ccaa0f59c3809eb1aed60e37b Mon Sep 17 00:00:00 2001 From: scxue Date: Thu, 8 May 2025 18:00:45 +0800 Subject: [PATCH 06/10] Add import for SanaVanillaAttnProcessor --- examples/research_projects/sana/train_sana_sprint_diffusers.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/examples/research_projects/sana/train_sana_sprint_diffusers.py b/examples/research_projects/sana/train_sana_sprint_diffusers.py index 2503e78636bf..0aea2c2cd3bd 100644 --- a/examples/research_projects/sana/train_sana_sprint_diffusers.py +++ b/examples/research_projects/sana/train_sana_sprint_diffusers.py @@ -21,7 +21,7 @@ import os import shutil from pathlib import Path -from typing import Callable +from typing import Callable, Optional import accelerate import numpy as np @@ -51,6 +51,7 @@ SanaSprintPipeline, SanaTransformer2DModel, ) +from diffusers.models.attention_processor import Attention from diffusers.optimization import get_scheduler from diffusers.training_utils import ( free_memory, From 04e1b02fb2221518ea50f6afb505e6fb43404e1e Mon Sep 17 00:00:00 2001 From: scxue Date: Thu, 8 May 2025 19:58:45 +0800 Subject: [PATCH 07/10] Add README file. --- examples/research_projects/sana/README.md | 95 +++++++++++++++++++ .../sana/train_sana_sprint_diffusers.sh | 2 + 2 files changed, 97 insertions(+) create mode 100644 examples/research_projects/sana/README.md diff --git a/examples/research_projects/sana/README.md b/examples/research_projects/sana/README.md new file mode 100644 index 000000000000..46d4f5f23095 --- /dev/null +++ b/examples/research_projects/sana/README.md @@ -0,0 +1,95 @@ +# Training SANA Sprint Diffuser + +This README explains how to use the provided Bash script commands to download a pre-trained teacher diffuser model and train it on a specific dataset. + + +## Setup + +### 1. Define Your Local Paths + +Set a variable for your desired output directory. This directory will store the downloaded model and the training checkpoints/results. + +```bash +your_local_path='output' # Or any other path you prefer +mkdir -p $your_local_path # Create the directory if it doesn't exist +``` + +### 2. Download the Pre-trained Model + +Download the SANA Sprint teacher model from Hugging Face Hub. The script uses the 1.6B parameter model. + +```bash +huggingface-cli download Efficient-Large-Model/SANA_Sprint_1.6B_1024px_teacher_diffusers --local-dir $your_local_path/SANA_Sprint_1.6B_1024px_teacher_diffusers +``` + +*(Optional: You can also download the 0.6B model by replacing the model name: `Efficient-Large-Model/Sana_Sprint_0.6B_1024px_teacher_diffusers`)* + +### 3. Acquire the Dataset Shards + +The training script in this example uses specific `.parquet` shards from the `brivangl/midjourney-v6-llava` dataset instead of downloading the entire dataset automatically via `dataset_name`. + +The script specifically uses these three files: +* `data/train_000.parquet` +* `data/train_001.parquet` +* `data/train_002.parquet` + + + +You can either: + +Let the script download the dataset automatically during first run + +Or download it manually + +**Note:** The full `brivangl/midjourney-v6-llava` dataset is much larger and contains many more shards. This script example explicitly trains *only* on the three specified shards. + +## Usage + +Once the model is downloaded, you can run the training script. + +```bash + +your_local_path='output' # Ensure this variable is set + +python train_sana_sprint_diffusers.py \ + --pretrained_model_name_or_path=$your_local_path/SANA_Sprint_1.6B_1024px_teacher_diffusers \ + --output_dir=$your_local_path \ + --mixed_precision=bf16 \ + --resolution=1024 \ + --learning_rate=1e-6 \ + --max_train_steps=30000 \ + --dataloader_num_workers=8 \ + --dataset_name='brivangl/midjourney-v6-llava' \ + --file_path data/train_000.parquet data/train_001.parquet data/train_002.parquet \ + --checkpointing_steps=500 --checkpoints_total_limit=10 \ + --train_batch_size=1 \ + --gradient_accumulation_steps=1 \ + --seed=453645634 \ + --train_largest_timestep \ + --misaligned_pairs_D \ + --gradient_checkpointing \ + --resume_from_checkpoint="latest" \ +``` + +### Explanation of Parameters + +* `--pretrained_model_name_or_path`: Path to the downloaded pre-trained model directory. +* `--output_dir`: Directory where training logs, checkpoints, and the final model will be saved. +* `--mixed_precision`: Use BF16 mixed precision for training, which can save memory and speed up training on compatible hardware. +* `--resolution`: The image resolution used for training (1024x1024). +* `--learning_rate`: The learning rate for the optimizer. +* `--max_train_steps`: The total number of training steps to perform. +* `--dataloader_num_workers`: Number of worker processes for loading data. Increase for faster data loading if your CPU and disk can handle it. +* `--dataset_name`: The name of the dataset on Hugging Face Hub (`brivangl/midjourney-v6-llava`). +* `--file_path`: **Specifies the local paths to the dataset shards to be used for training.** In this case, `data/train_000.parquet`, `data/train_001.parquet`, and `data/train_002.parquet`. +* `--checkpointing_steps`: Save a training checkpoint every X steps. +* `--checkpoints_total_limit`: Maximum number of checkpoints to keep. Older checkpoints will be deleted. +* `--train_batch_size`: The batch size per GPU. +* `--gradient_accumulation_steps`: Number of steps to accumulate gradients before performing an optimizer step. +* `--seed`: Random seed for reproducibility. +* `--train_largest_timestep`: A specific training strategy focusing on larger timesteps. +* `--misaligned_pairs_D`: Another specific training strategy to add misaligned image-text pairs as fake data for GAN. +* `--gradient_checkpointing`: Enable gradient checkpointing to save GPU memory. +* `--resume_from_checkpoint`: Allows resuming training from the latest saved checkpoint in the `--output_dir`. + + diff --git a/examples/research_projects/sana/train_sana_sprint_diffusers.sh b/examples/research_projects/sana/train_sana_sprint_diffusers.sh index bf6f47f9d731..301fe5e429a5 100644 --- a/examples/research_projects/sana/train_sana_sprint_diffusers.sh +++ b/examples/research_projects/sana/train_sana_sprint_diffusers.sh @@ -2,6 +2,8 @@ your_local_path='output' huggingface-cli download Efficient-Large-Model/SANA_Sprint_1.6B_1024px_teacher_diffusers --local-dir $your_local_path/SANA_Sprint_1.6B_1024px_teacher_diffusers +# or Sana_Sprint_0.6B_1024px_teacher_diffusers + python train_sana_sprint_diffusers.py \ --pretrained_model_name_or_path=$your_local_path/SANA_Sprint_1.6B_1024px_teacher_diffusers \ --output_dir=$your_local_path \ From 5951f8fee82b34582966b92612b3978ebbf4823c Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Thu, 8 May 2025 17:46:19 +0530 Subject: [PATCH 08/10] Apply suggestions from code review --- examples/research_projects/sana/README.md | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/examples/research_projects/sana/README.md b/examples/research_projects/sana/README.md index 46d4f5f23095..7bf67f1c4b53 100644 --- a/examples/research_projects/sana/README.md +++ b/examples/research_projects/sana/README.md @@ -1,11 +1,11 @@ # Training SANA Sprint Diffuser -This README explains how to use the provided Bash script commands to download a pre-trained teacher diffuser model and train it on a specific dataset. +This README explains how to use the provided bash script commands to download a pre-trained teacher diffuser model and train it on a specific dataset, following the [SANA Sprint methodology](https://arxiv.org/abs/2503.09641). ## Setup -### 1. Define Your Local Paths +### 1. Define the local paths Set a variable for your desired output directory. This directory will store the downloaded model and the training checkpoints/results. @@ -14,7 +14,7 @@ your_local_path='output' # Or any other path you prefer mkdir -p $your_local_path # Create the directory if it doesn't exist ``` -### 2. Download the Pre-trained Model +### 2. Download the pre-trained model Download the SANA Sprint teacher model from Hugging Face Hub. The script uses the 1.6B parameter model. @@ -24,7 +24,7 @@ huggingface-cli download Efficient-Large-Model/SANA_Sprint_1.6B_1024px_teacher_d *(Optional: You can also download the 0.6B model by replacing the model name: `Efficient-Large-Model/Sana_Sprint_0.6B_1024px_teacher_diffusers`)* -### 3. Acquire the Dataset Shards +### 3. Acquire the dataset shards The training script in this example uses specific `.parquet` shards from the `brivangl/midjourney-v6-llava` dataset instead of downloading the entire dataset automatically via `dataset_name`. @@ -71,7 +71,7 @@ python train_sana_sprint_diffusers.py \ --resume_from_checkpoint="latest" \ ``` -### Explanation of Parameters +### Explanation of parameters * `--pretrained_model_name_or_path`: Path to the downloaded pre-trained model directory. * `--output_dir`: Directory where training logs, checkpoints, and the final model will be saved. From 566aa64d2c4859081f9c29685998fb67bb0c5da7 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Thu, 8 May 2025 17:52:23 +0530 Subject: [PATCH 09/10] style --- .../sana/train_sana_sprint_diffusers.py | 246 +++++++++++------- 1 file changed, 153 insertions(+), 93 deletions(-) diff --git a/examples/research_projects/sana/train_sana_sprint_diffusers.py b/examples/research_projects/sana/train_sana_sprint_diffusers.py index 0aea2c2cd3bd..335d9c377c06 100644 --- a/examples/research_projects/sana/train_sana_sprint_diffusers.py +++ b/examples/research_projects/sana/train_sana_sprint_diffusers.py @@ -87,6 +87,7 @@ "User Prompt: ", ] + class SanaVanillaAttnProcessor: r""" Processor for implementing scaled dot-product attention to support JVP calculation during training. @@ -96,7 +97,8 @@ def __init__(self): pass @staticmethod - def scaled_dot_product_attention(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, scale=None + def scaled_dot_product_attention( + query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, scale=None ) -> torch.Tensor: B, H, L, S = *query.size()[:-1], key.size(-2) scale_factor = 1 / math.sqrt(query.size(-1)) if scale is None else scale @@ -173,44 +175,42 @@ class Text2ImageDataset(Dataset): """ A PyTorch Dataset class for loading text-image pairs from a HuggingFace dataset. This dataset is designed for text-to-image generation tasks. - Args: - hf_dataset (datasets.Dataset): A HuggingFace dataset containing 'image' (bytes) - and 'llava' (text) fields. Note that 'llava' is the field name for text descriptions in this specific dataset - you may need to adjust this key if using a different HuggingFace dataset with a different text field name. - resolution (int, optional): Target resolution for image resizing. Defaults to 1024. - + hf_dataset (datasets.Dataset): + A HuggingFace dataset containing 'image' (bytes) and 'llava' (text) fields. Note that 'llava' is the field name for text descriptions in this specific dataset - you may need to adjust this key if using a different HuggingFace dataset with a different text field name. + resolution (int, optional): Target resolution for image resizing. Defaults to 1024. Returns: dict: A dictionary containing: - 'text': The text description (str) - 'image': The processed image tensor (torch.Tensor) of shape [3, resolution, resolution] """ + def __init__(self, hf_dataset, resolution=1024): self.dataset = hf_dataset - self.transform = T.Compose([ - T.Lambda(lambda img: img.convert("RGB")), - T.Resize(resolution), # Image.BICUBIC - T.CenterCrop(resolution), - T.ToTensor(), - T.Normalize([0.5], [0.5]), - ]) + self.transform = T.Compose( + [ + T.Lambda(lambda img: img.convert("RGB")), + T.Resize(resolution), # Image.BICUBIC + T.CenterCrop(resolution), + T.ToTensor(), + T.Normalize([0.5], [0.5]), + ] + ) def __len__(self): return len(self.dataset) def __getitem__(self, idx): item = self.dataset[idx] - text = item['llava'] - image_bytes = item['image'] + text = item["llava"] + image_bytes = item["image"] # Convert bytes to PIL Image image = Image.open(io.BytesIO(image_bytes)) image_tensor = self.transform(image) - return { - 'text': text, - 'image': image_tensor - } + return {"text": text, "image": image_tensor} def save_model_card( @@ -389,8 +389,7 @@ def parse_args(input_args=None): ) parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.") # ----Image Processing---- - parser.add_argument("--file_path", nargs='+', required=True, - help="List of parquet files (space-separated)") + parser.add_argument("--file_path", nargs="+", required=True, help="List of parquet files (space-separated)") parser.add_argument( "--dataset_name", type=str, @@ -524,21 +523,45 @@ def parse_args(input_args=None): parser.add_argument( "--logit_std", type=float, default=1.6, help="std to use when using the `'logit_normal'` weighting scheme." ) - parser.add_argument("--logit_mean_discriminator", type=float, default=-0.6, help="Logit mean for discriminator timestep sampling") - parser.add_argument("--logit_std_discriminator", type=float, default=1.0, help="Logit std for discriminator timestep sampling") + parser.add_argument( + "--logit_mean_discriminator", type=float, default=-0.6, help="Logit mean for discriminator timestep sampling" + ) + parser.add_argument( + "--logit_std_discriminator", type=float, default=1.0, help="Logit std for discriminator timestep sampling" + ) parser.add_argument("--ladd_multi_scale", action="store_true", help="Whether to use multi-scale discriminator") - parser.add_argument("--head_block_ids", type=int, nargs="+", default=[2, 8, 14, 19], help="Specify which transformer blocks to use for discriminator heads") + parser.add_argument( + "--head_block_ids", + type=int, + nargs="+", + default=[2, 8, 14, 19], + help="Specify which transformer blocks to use for discriminator heads", + ) parser.add_argument("--adv_lambda", type=float, default=0.5, help="Weighting coefficient for adversarial loss") parser.add_argument("--scm_lambda", type=float, default=1.0, help="Weighting coefficient for SCM loss") parser.add_argument("--gradient_clip", type=float, default=0.1, help="Threshold for gradient clipping") - parser.add_argument("--sigma_data", type=float, default=0.5, help="Standard deviation of data distribution is supposed to be 0.5") - parser.add_argument("--tangent_warmup_steps", type=int, default=4000, help="Number of warmup steps for tangent vectors") - parser.add_argument("--guidance_embeds_scale", type=float, default=0.1, help="Scaling factor for guidance embeddings") - parser.add_argument("--scm_cfg_scale", type=float, nargs="+", default=[4, 4.5, 5], help="Range for classifier-free guidance scale") - parser.add_argument("--train_largest_timestep", action="store_true", help="Whether to enable special training for large timesteps") + parser.add_argument( + "--sigma_data", type=float, default=0.5, help="Standard deviation of data distribution is supposed to be 0.5" + ) + parser.add_argument( + "--tangent_warmup_steps", type=int, default=4000, help="Number of warmup steps for tangent vectors" + ) + parser.add_argument( + "--guidance_embeds_scale", type=float, default=0.1, help="Scaling factor for guidance embeddings" + ) + parser.add_argument( + "--scm_cfg_scale", type=float, nargs="+", default=[4, 4.5, 5], help="Range for classifier-free guidance scale" + ) + parser.add_argument( + "--train_largest_timestep", action="store_true", help="Whether to enable special training for large timesteps" + ) parser.add_argument("--largest_timestep", type=float, default=1.57080, help="Maximum timestep value") - parser.add_argument("--largest_timestep_prob", type=float, default=0.5, help="Sampling probability for large timesteps") - parser.add_argument("--misaligned_pairs_D", action="store_true", help="Add misaligned sample pairs for discriminator") + parser.add_argument( + "--largest_timestep_prob", type=float, default=0.5, help="Sampling probability for large timesteps" + ) + parser.add_argument( + "--misaligned_pairs_D", action="store_true", help="Add misaligned sample pairs for discriminator" + ) parser.add_argument( "--optimizer", type=str, @@ -663,12 +686,10 @@ def parse_args(input_args=None): else: 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 - return args @@ -791,7 +812,13 @@ def get_features(module, input, output): if i in self.block_hooks: hooks.append(block.register_forward_hook(get_features)) - self.transformer(hidden_states=hidden_states, timestep=timestep, encoder_hidden_states=encoder_hidden_states, return_logvar=False, **kwargs) + self.transformer( + hidden_states=hidden_states, + timestep=timestep, + encoder_hidden_states=encoder_hidden_states, + return_logvar=False, + **kwargs, + ) for hook in hooks: hook.remove() @@ -825,6 +852,7 @@ def state_dict(self): def __getattr__(self, name): return getattr(self.disc, name) + class SanaTrigFlow(SanaTransformer2DModel): def __init__(self, original_model, guidance=False): self.__dict__ = original_model.__dict__ @@ -836,7 +864,9 @@ def __init__(self, original_model, guidance=False): torch.nn.init.xavier_uniform_(self.logvar_linear.weight) torch.nn.init.constant_(self.logvar_linear.bias, 0) - def forward(self, hidden_states, encoder_hidden_states, timestep, guidance=None, jvp=False, return_logvar=False, **kwargs): + def forward( + self, hidden_states, encoder_hidden_states, timestep, guidance=None, jvp=False, return_logvar=False, **kwargs + ): batch_size = hidden_states.shape[0] latents = hidden_states prompt_embeds = encoder_hidden_states @@ -858,17 +888,28 @@ def forward(self, hidden_states, encoder_hidden_states, timestep, guidance=None, if jvp and self.gradient_checkpointing: self.gradient_checkpointing = False - model_out = super().forward(hidden_states=latent_model_input, encoder_hidden_states=prompt_embeds, timestep=flow_timestep, guidance=guidance, **kwargs)[0] + model_out = super().forward( + hidden_states=latent_model_input, + encoder_hidden_states=prompt_embeds, + timestep=flow_timestep, + guidance=guidance, + **kwargs, + )[0] self.gradient_checkpointing = True else: - model_out = super().forward(hidden_states=latent_model_input, encoder_hidden_states=prompt_embeds, timestep=flow_timestep, guidance=guidance, **kwargs)[0] - + model_out = super().forward( + hidden_states=latent_model_input, + encoder_hidden_states=prompt_embeds, + timestep=flow_timestep, + guidance=guidance, + **kwargs, + )[0] # Flow --> TrigFlow Transformation - trigflow_model_out = ((1 - 2 * flow_timestep_expanded) * latent_model_input + (1 - 2 * flow_timestep_expanded + 2 * flow_timestep_expanded**2) * model_out) / torch.sqrt( - flow_timestep_expanded**2 + (1 - flow_timestep_expanded) ** 2 - ) - + trigflow_model_out = ( + (1 - 2 * flow_timestep_expanded) * latent_model_input + + (1 - 2 * flow_timestep_expanded + 2 * flow_timestep_expanded**2) * model_out + ) / torch.sqrt(flow_timestep_expanded**2 + (1 - flow_timestep_expanded) ** 2) if self.guidance and guidance is not None: timestep, embedded_timestep = self.time_embed( @@ -883,15 +924,11 @@ def forward(self, hidden_states, encoder_hidden_states, timestep, guidance=None, logvar = self.logvar_linear(embedded_timestep) return trigflow_model_out, logvar - return (trigflow_model_out,) -def compute_density_for_timestep_sampling_scm( - batch_size: int, logit_mean: float = None, logit_std: float = None -): - """Compute the density for sampling the timesteps when doing Sana-Sprint training. - """ +def compute_density_for_timestep_sampling_scm(batch_size: int, logit_mean: float = None, logit_std: float = None): + """Compute the density for sampling the timesteps when doing Sana-Sprint training.""" sigma = torch.randn(batch_size, device="cpu") sigma = (sigma * logit_std + logit_mean).exp() u = torch.atan(sigma / 0.5) # TODO: 0.5 should be a hyper-parameter @@ -980,23 +1017,31 @@ def main(args): ) ori_transformer = SanaTransformer2DModel.from_pretrained( - args.pretrained_model_name_or_path, subfolder="transformer", revision=args.revision, variant=args.variant, + args.pretrained_model_name_or_path, + subfolder="transformer", + revision=args.revision, + variant=args.variant, guidance_embeds=True, ) ori_transformer.set_attn_processor(SanaVanillaAttnProcessor()) ori_transformer_no_guide = SanaTransformer2DModel.from_pretrained( - args.pretrained_model_name_or_path, subfolder="transformer", revision=args.revision, variant=args.variant, - guidance_embeds=False + args.pretrained_model_name_or_path, + subfolder="transformer", + revision=args.revision, + variant=args.variant, + guidance_embeds=False, ) - original_state_dict = load_file(f"{args.pretrained_model_name_or_path}/transformer/diffusion_pytorch_model.safetensors") + original_state_dict = load_file( + f"{args.pretrained_model_name_or_path}/transformer/diffusion_pytorch_model.safetensors" + ) param_mapping = { - 'time_embed.emb.timestep_embedder.linear_1.weight': 'time_embed.timestep_embedder.linear_1.weight', - 'time_embed.emb.timestep_embedder.linear_1.bias': 'time_embed.timestep_embedder.linear_1.bias', - 'time_embed.emb.timestep_embedder.linear_2.weight': 'time_embed.timestep_embedder.linear_2.weight', - 'time_embed.emb.timestep_embedder.linear_2.bias': 'time_embed.timestep_embedder.linear_2.bias', + "time_embed.emb.timestep_embedder.linear_1.weight": "time_embed.timestep_embedder.linear_1.weight", + "time_embed.emb.timestep_embedder.linear_1.bias": "time_embed.timestep_embedder.linear_1.bias", + "time_embed.emb.timestep_embedder.linear_2.weight": "time_embed.timestep_embedder.linear_2.weight", + "time_embed.emb.timestep_embedder.linear_2.bias": "time_embed.timestep_embedder.linear_2.bias", } for src_key, dst_key in param_mapping.items(): @@ -1009,13 +1054,13 @@ def main(args): target_device = accelerator.device param_w1 = guidance_embedder_module.linear_1.weight - zero_state_dict['linear_1.weight'] = torch.zeros(param_w1.shape, device=target_device) + zero_state_dict["linear_1.weight"] = torch.zeros(param_w1.shape, device=target_device) param_b1 = guidance_embedder_module.linear_1.bias - zero_state_dict['linear_1.bias'] = torch.zeros(param_b1.shape, device=target_device) + zero_state_dict["linear_1.bias"] = torch.zeros(param_b1.shape, device=target_device) param_w2 = guidance_embedder_module.linear_2.weight - zero_state_dict['linear_2.weight'] = torch.zeros(param_w2.shape, device=target_device) + zero_state_dict["linear_2.weight"] = torch.zeros(param_w2.shape, device=target_device) param_b2 = guidance_embedder_module.linear_2.bias - zero_state_dict['linear_2.bias'] = torch.zeros(param_b2.shape, device=target_device) + zero_state_dict["linear_2.bias"] = torch.zeros(param_b2.shape, device=target_device) guidance_embedder_module.load_state_dict(zero_state_dict, strict=False, assign=True) transformer = SanaTrigFlow(ori_transformer, guidance=True).train() @@ -1098,10 +1143,7 @@ def save_model_hook(models, weights, output_dir): # Handle discriminator model (only save heads) elif isinstance(unwrapped_model, type(unwrap_model(disc))): # Save only the heads - torch.save( - unwrapped_model.heads.state_dict(), - os.path.join(output_dir, "disc_heads.pt") - ) + torch.save(unwrapped_model.heads.state_dict(), os.path.join(output_dir, "disc_heads.pt")) else: raise ValueError(f"unexpected save model: {unwrapped_model.__class__}") @@ -1131,7 +1173,7 @@ def load_model_hook(models, input_dir): else: # DeepSpeed case transformer_ = SanaTransformer2DModel.from_pretrained(input_dir, subfolder="transformer") # noqa: F841 - disc_heads_state_dict = torch.load(os.path.join(input_dir, "disc_heads.pt")) + disc_heads_state_dict = torch.load(os.path.join(input_dir, "disc_heads.pt")) # noqa: F841 # You'll need to handle how to load the heads in DeepSpeed case accelerator.register_save_state_pre_hook(save_model_hook) @@ -1180,7 +1222,7 @@ def load_model_hook(models, input_dir): hf_dataset = load_dataset( args.dataset_name, data_files=args.file_path, - split='train', + split="train", ) train_dataset = Text2ImageDataset( @@ -1194,7 +1236,7 @@ def load_model_hook(models, input_dir): num_workers=args.dataloader_num_workers, pin_memory=True, persistent_workers=True, - shuffle=True + shuffle=True, ) # Scheduler and math around the number of training steps. @@ -1214,8 +1256,10 @@ def load_model_hook(models, input_dir): ) # Prepare everything with our `accelerator`. - transformer, pretrained_model, disc, optimizer_G, optimizer_D, train_dataloader, lr_scheduler = accelerator.prepare( - transformer, pretrained_model, disc, optimizer_G, optimizer_D, train_dataloader, lr_scheduler + transformer, pretrained_model, disc, optimizer_G, optimizer_D, train_dataloader, lr_scheduler = ( + accelerator.prepare( + transformer, pretrained_model, disc, optimizer_G, optimizer_D, train_dataloader, lr_scheduler + ) ) # We need to recalculate our total training steps as the size of the training dataloader may have changed. @@ -1230,8 +1274,7 @@ def load_model_hook(models, input_dir): if accelerator.is_main_process: tracker_name = "sana-sprint" config = { - k: str(v) if not isinstance(v, (int, float, str, bool, torch.Tensor)) else v - for k, v in vars(args).items() + k: str(v) if not isinstance(v, (int, float, str, bool, torch.Tensor)) else v for k, v in vars(args).items() } accelerator.init_trackers(tracker_name, config=config) @@ -1285,7 +1328,7 @@ def load_model_hook(models, input_dir): disable=not accelerator.is_local_main_process, ) - phase = 'G' + phase = "G" vae_config_scaling_factor = vae.config.scaling_factor sigma_data = args.sigma_data negative_prompt = [""] * args.train_batch_size @@ -1326,13 +1369,11 @@ def load_model_hook(models, input_dir): logit_std=args.logit_std, ).to(accelerator.device) - # Add noise according to TrigFlow. # zt = cos(t) * x + sin(t) * noise t = u.view(-1, 1, 1, 1) noisy_model_input = torch.cos(t) * model_input + torch.sin(t) * noise - scm_cfg_scale = torch.tensor( np.random.choice(args.scm_cfg_scale, size=bsz, replace=True), device=accelerator.device, @@ -1340,7 +1381,13 @@ def load_model_hook(models, input_dir): def model_wrapper(scaled_x_t, t): pred, logvar = accelerator.unwrap_model(transformer)( - hidden_states=scaled_x_t, timestep=t.flatten(), encoder_hidden_states=prompt_embeds, encoder_attention_mask=prompt_attention_mask, guidance=(scm_cfg_scale.flatten() * args.guidance_embeds_scale), jvp=True, return_logvar=True + hidden_states=scaled_x_t, + timestep=t.flatten(), + encoder_hidden_states=prompt_embeds, + encoder_attention_mask=prompt_attention_mask, + guidance=(scm_cfg_scale.flatten() * args.guidance_embeds_scale), + jvp=True, + return_logvar=True, ) return pred, logvar @@ -1355,9 +1402,11 @@ def model_wrapper(scaled_x_t, t): cfg_y = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) cfg_y_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0) - cfg_pretrain_pred = pretrained_model( - hidden_states=(cfg_x_t / sigma_data), timestep=cfg_t.flatten(), encoder_hidden_states=cfg_y, encoder_attention_mask=cfg_y_mask + hidden_states=(cfg_x_t / sigma_data), + timestep=cfg_t.flatten(), + encoder_hidden_states=cfg_y, + encoder_attention_mask=cfg_y_mask, )[0] cfg_dxt_dt = sigma_data * cfg_pretrain_pred @@ -1367,28 +1416,24 @@ def model_wrapper(scaled_x_t, t): scm_cfg_scale = scm_cfg_scale.view(-1, 1, 1, 1) dxt_dt = dxt_dt_uncond + scm_cfg_scale * (dxt_dt - dxt_dt_uncond) - v_x = torch.cos(t) * torch.sin(t) * dxt_dt / sigma_data v_t = torch.cos(t) * torch.sin(t) - # Adapt from https://github.com/xandergos/sCM-mnist/blob/master/train_consistency.py with torch.no_grad(): F_theta, F_theta_grad, logvar = torch.func.jvp( model_wrapper, (noisy_model_input / sigma_data, t), (v_x, v_t), has_aux=True ) - F_theta, logvar = transformer( - hidden_states = (noisy_model_input / sigma_data), + hidden_states=(noisy_model_input / sigma_data), timestep=t.flatten(), encoder_hidden_states=prompt_embeds, encoder_attention_mask=prompt_attention_mask, guidance=(scm_cfg_scale.flatten() * args.guidance_embeds_scale), - return_logvar=True + return_logvar=True, ) - logvar = logvar.view(-1, 1, 1, 1) F_theta_grad = F_theta_grad.detach() F_theta_minus = F_theta.detach() @@ -1417,10 +1462,8 @@ def model_wrapper(scaled_x_t, t): loss_no_logvar = weight * torch.square(F_theta - F_theta_minus - g) loss_no_logvar = loss_no_logvar.mean() - loss_no_weight = l2_loss.mean() g_norm = g_norm.mean() - pred_x_0 = torch.cos(t) * noisy_model_input - torch.sin(t) * F_theta * sigma_data if args.train_largest_timestep: @@ -1439,7 +1482,7 @@ def model_wrapper(scaled_x_t, t): x_t_new = torch.cos(t_new) * model_input + torch.sin(t_new) * z_new F_theta = transformer( - hidden_states = (x_t_new / sigma_data), + hidden_states=(x_t_new / sigma_data), timestep=t_new.flatten(), encoder_hidden_states=prompt_embeds, encoder_attention_mask=prompt_attention_mask, @@ -1448,7 +1491,6 @@ def model_wrapper(scaled_x_t, t): jvp=False, )[0] - pred_x_0 = torch.cos(t_new) * x_t_new - torch.sin(t_new) * F_theta * sigma_data # Sample timesteps for discriminator @@ -1464,7 +1506,12 @@ def model_wrapper(scaled_x_t, t): noised_predicted_x0 = torch.cos(t_D) * pred_x_0 + torch.sin(t_D) * z_D # Calculate adversarial loss - pred_fake = disc(hidden_states=(noised_predicted_x0 / sigma_data), timestep=t_D.flatten(), encoder_hidden_states=prompt_embeds, encoder_attention_mask=prompt_attention_mask) + pred_fake = disc( + hidden_states=(noised_predicted_x0 / sigma_data), + timestep=t_D.flatten(), + encoder_hidden_states=prompt_embeds, + encoder_attention_mask=prompt_attention_mask, + ) adv_loss = -torch.mean(pred_fake) # Total loss = sCM loss + LADD loss @@ -1509,7 +1556,7 @@ def model_wrapper(scaled_x_t, t): noisy_model_input = torch.cos(t) * model_input + torch.sin(t) * z_new # here F_theta = transformer( - hidden_states = (noisy_model_input / sigma_data), + hidden_states=(noisy_model_input / sigma_data), timestep=t.flatten(), encoder_hidden_states=prompt_embeds, encoder_attention_mask=prompt_attention_mask, @@ -1557,11 +1604,20 @@ def model_wrapper(scaled_x_t, t): prompt_embeds = torch.cat([prompt_embeds, prompt_embeds], dim=0) prompt_attention_mask = torch.cat([prompt_attention_mask, prompt_attention_mask], dim=0) - # Calculate D loss - pred_fake = disc(hidden_states=(noised_predicted_x0 / sigma_data), timestep=t_D_fake.flatten(), encoder_hidden_states=prompt_embeds, encoder_attention_mask=prompt_attention_mask) - pred_true = disc(hidden_states=(noised_latents / sigma_data), timestep=t_D_real.flatten(), encoder_hidden_states=prompt_embeds, encoder_attention_mask=prompt_attention_mask) + pred_fake = disc( + hidden_states=(noised_predicted_x0 / sigma_data), + timestep=t_D_fake.flatten(), + encoder_hidden_states=prompt_embeds, + encoder_attention_mask=prompt_attention_mask, + ) + pred_true = disc( + hidden_states=(noised_latents / sigma_data), + timestep=t_D_real.flatten(), + encoder_hidden_states=prompt_embeds, + encoder_attention_mask=prompt_attention_mask, + ) # hinge loss loss_real = torch.mean(F.relu(1.0 - pred_true)) @@ -1617,7 +1673,11 @@ def model_wrapper(scaled_x_t, t): accelerator.save_state(save_path) logger.info(f"Saved state to {save_path}") - logs = {"scm_loss": loss.detach().item(), "adv_loss": adv_loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]} + logs = { + "scm_loss": loss.detach().item(), + "adv_loss": adv_loss.detach().item(), + "lr": lr_scheduler.get_last_lr()[0], + } progress_bar.set_postfix(**logs) accelerator.log(logs, step=global_step) From 740baa9e1fbed0535d2723c1e70fe44b62fad680 Mon Sep 17 00:00:00 2001 From: Junsong Chen Date: Thu, 8 May 2025 20:52:37 +0800 Subject: [PATCH 10/10] Update examples/research_projects/sana/README.md --- examples/research_projects/sana/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/research_projects/sana/README.md b/examples/research_projects/sana/README.md index 7bf67f1c4b53..1dbae667c8f9 100644 --- a/examples/research_projects/sana/README.md +++ b/examples/research_projects/sana/README.md @@ -26,7 +26,7 @@ huggingface-cli download Efficient-Large-Model/SANA_Sprint_1.6B_1024px_teacher_d ### 3. Acquire the dataset shards -The training script in this example uses specific `.parquet` shards from the `brivangl/midjourney-v6-llava` dataset instead of downloading the entire dataset automatically via `dataset_name`. +The training script in this example uses specific `.parquet` shards from a randomly selected `brivangl/midjourney-v6-llava` dataset instead of downloading the entire dataset automatically via `dataset_name`. The script specifically uses these three files: * `data/train_000.parquet`