"""

random.seed(seed)

np.random.seed(seed)

torch.manual_seed(seed)

if is_torch_npu_available():

torch.npu.manual_seed_all(seed)

else:

torch.cuda.manual_seed_all(seed)

"""

alphas_cumprod = noise_scheduler.alphas_cumprod

sqrt_alphas_cumprod = alphas_cumprod**0.5

sqrt_one_minus_alphas_cumprod = (1.0 - alphas_cumprod) ** 0.5

# Expand the tensors.

# Adapted from https://github.com/TiankaiHang/Min-SNR-Diffusion-Training/blob/521b624bd70c67cee4bdf49225915f5945a872e3/guided_diffusion/gaussian_diffusion.py#L1026

sqrt_alphas_cumprod = sqrt_alphas_cumprod.to(device=timesteps.device)[timesteps].float()

while len(sqrt_alphas_cumprod.shape) < len(timesteps.shape):

sqrt_alphas_cumprod = sqrt_alphas_cumprod[..., None]

alpha = sqrt_alphas_cumprod.expand(timesteps.shape)

sqrt_one_minus_alphas_cumprod = sqrt_one_minus_alphas_cumprod.to(device=timesteps.device)[timesteps].float()

while len(sqrt_one_minus_alphas_cumprod.shape) < len(timesteps.shape):

sqrt_one_minus_alphas_cumprod = sqrt_one_minus_alphas_cumprod[..., None]

sigma = sqrt_one_minus_alphas_cumprod.expand(timesteps.shape)

# Compute SNR.

snr = (alpha / sigma) ** 2

logits: torch.Tensor,

labels: torch.Tensor,

*,

lambda_conf: float = 0.0,

temperature: float = 1.0,

per_token_weights: torch.Tensor | None = None,

ignore_index: int = -100,

) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:

"""

if logits.ndim < 2:

raise ValueError(f"`logits` must have at least 2 dims, got shape {tuple(logits.shape)}.")

if labels.shape != logits.shape[:-1]:

raise ValueError(

f"`labels` shape must match `logits.shape[:-1]`, got labels={tuple(labels.shape)} logits={tuple(logits.shape)}."

)

if temperature <= 0:

raise ValueError(f"`temperature` must be > 0, got {temperature}.")

valid = labels.ne(ignore_index)

if per_token_weights is None:

weights = torch.ones_like(labels, dtype=logits.dtype)

else:

if per_token_weights.shape != labels.shape:

raise ValueError(

f"`per_token_weights` shape must match `labels` shape, got {tuple(per_token_weights.shape)} != {tuple(labels.shape)}."

)

weights = per_token_weights.to(dtype=logits.dtype)

# Supervised CE (optionally weighted).

vocab_size = logits.shape[-1]

per_token_nll = F.cross_entropy(

logits.reshape(-1, vocab_size),

labels.reshape(-1),

reduction="none",

ignore_index=ignore_index,

).reshape_as(labels)

denom_sft = (weights * valid.to(weights.dtype)).sum().clamp_min(1)

loss_sft = (per_token_nll * weights * valid.to(per_token_nll.dtype)).sum() / denom_sft

# Confidence loss: penalize entropy only where prediction is already correct.

if lambda_conf == 0.0:

loss_conf = torch.zeros((), device=logits.device, dtype=loss_sft.dtype)

return loss_sft, loss_sft, loss_conf

with torch.no_grad():

pred = logits.argmax(dim=-1)

correct = valid & pred.eq(labels)

scaled_logits = logits.float()

if temperature != 1.0:

scaled_logits = scaled_logits / float(temperature)

probs = torch.softmax(scaled_logits, dim=-1)

eps = torch.finfo(probs.dtype).tiny

log_probs = torch.log(probs.clamp_min(eps))

entropy = -(probs * log_probs).sum(dim=-1).to(dtype=logits.dtype)

denom_conf = (weights * correct.to(weights.dtype)).sum().clamp_min(1)

loss_conf = (entropy * weights * correct.to(entropy.dtype)).sum() / denom_conf

loss = loss_sft + float(lambda_conf) * loss_conf

"""

if not is_torchvision_available():

raise ImportError(

"Please make sure to install `torchvision` to be able to use the `resolve_interpolation_mode()` function."

)

if interpolation_type == "bilinear":

interpolation_mode = transforms.InterpolationMode.BILINEAR

elif interpolation_type == "bicubic":

interpolation_mode = transforms.InterpolationMode.BICUBIC

elif interpolation_type == "box":

interpolation_mode = transforms.InterpolationMode.BOX

elif interpolation_type == "nearest":

interpolation_mode = transforms.InterpolationMode.NEAREST

elif interpolation_type == "nearest_exact":

interpolation_mode = transforms.InterpolationMode.NEAREST_EXACT

elif interpolation_type == "hamming":

interpolation_mode = transforms.InterpolationMode.HAMMING

elif interpolation_type == "lanczos":

interpolation_mode = transforms.InterpolationMode.LANCZOS

else:

raise ValueError(

f"The given interpolation mode {interpolation_type} is not supported. Currently supported interpolation"

f" modes are `bilinear`, `bicubic`, `box`, `nearest`, `nearest_exact`, `hamming`, and `lanczos`."

)

unet: UNet2DConditionModel,

noise_scheduler: SchedulerMixin,

timesteps: torch.Tensor,

noise: torch.Tensor,

noisy_latents: torch.Tensor,

target: torch.Tensor,

encoder_hidden_states: torch.Tensor,

dream_detail_preservation: float = 1.0,

) -> tuple[torch.Tensor, torch.Tensor]:

"""

alphas_cumprod = noise_scheduler.alphas_cumprod.to(timesteps.device)[timesteps, None, None, None]

sqrt_one_minus_alphas_cumprod = (1.0 - alphas_cumprod) ** 0.5

# The paper uses lambda = sqrt(1 - alpha) ** p, with p = 1 in their experiments.

dream_lambda = sqrt_one_minus_alphas_cumprod**dream_detail_preservation

pred = None

with torch.no_grad():

pred = unet(noisy_latents, timesteps, encoder_hidden_states).sample

_noisy_latents, _target = (None, None)

if noise_scheduler.config.prediction_type == "epsilon":

predicted_noise = pred

delta_noise = (noise - predicted_noise).detach()

delta_noise.mul_(dream_lambda)

_noisy_latents = noisy_latents.add(sqrt_one_minus_alphas_cumprod * delta_noise)

_target = target.add(delta_noise)

elif noise_scheduler.config.prediction_type == "v_prediction":

raise NotImplementedError("DREAM has not been implemented for v-prediction")

else:

raise ValueError(f"Unknown prediction type {noise_scheduler.config.prediction_type}")

r"""

lora_state_dict = {}

for name, module in unet.named_modules():

if hasattr(module, "set_lora_layer"):

lora_layer = getattr(module, "lora_layer")

if lora_layer is not None:

current_lora_layer_sd = lora_layer.state_dict()

for lora_layer_matrix_name, lora_param in current_lora_layer_sd.items():

# The matrix name can either be "down" or "up".

lora_state_dict[f"{name}.lora.{lora_layer_matrix_name}"] = lora_param

"""

if not isinstance(model, list):

model = [model]

for m in model:

for param in m.parameters():

# only upcast trainable parameters into fp32

if param.requires_grad:

lora_state_dict: dict[str, torch.Tensor], prefix: str, text_encoder: torch.nn.Module

):

"""

text_encoder_state_dict = {

f"{k.replace(prefix, '')}": v for k, v in lora_state_dict.items() if k.startswith(prefix)

}

text_encoder_state_dict = convert_state_dict_to_peft(convert_state_dict_to_diffusers(text_encoder_state_dict))

metadatas = {}

for module_name, module in modules_to_save.items():

if module is not None:

metadatas[f"{module_name}_lora_adapter_metadata"] = module.peft_config["default"].to_dict()

weighting_scheme: str,

batch_size: int,

logit_mean: float = None,

logit_std: float = None,

mode_scale: float = None,

device: torch.device | str = "cpu",

generator: torch.Generator | None = None,

):

"""

if weighting_scheme == "logit_normal":

u = torch.normal(mean=logit_mean, std=logit_std, size=(batch_size,), device=device, generator=generator)

u = torch.nn.functional.sigmoid(u)

elif weighting_scheme == "mode":

u = torch.rand(size=(batch_size,), device=device, generator=generator)

u = 1 - u - mode_scale * (torch.cos(math.pi * u / 2) ** 2 - 1 + u)

else:

u = torch.rand(size=(batch_size,), device=device, generator=generator)

"""

if weighting_scheme == "sigma_sqrt":

weighting = (sigmas**-2.0).float()

elif weighting_scheme == "cosmap":

bot = 1 - 2 * sigmas + 2 * sigmas**2

weighting = 2 / (math.pi * bot)

else:

weighting = torch.ones_like(sigmas)

"""

gc.collect()

if torch.cuda.is_available():

torch.cuda.empty_cache()

elif torch.backends.mps.is_available():

torch.mps.empty_cache()

elif is_torch_npu_available():

torch_npu.npu.empty_cache()

elif hasattr(torch, "xpu") and torch.xpu.is_available():

"""

if offload:

is_model = not any(isinstance(m, DiffusionPipeline) for m in modules)

# record where each module was

if is_model:

original_devices = [next(m.parameters()).device for m in modules]

else:

assert len(modules) == 1

# For DiffusionPipeline, wrap the device in a list to make it iterable

original_devices = [modules[0].device]

# move to target device

for m in modules:

m.to(device)

try:

yield

finally:

if offload:

# move back to original devices

for m, orig_dev in zip(modules, original_devices):

"""Parses a string defining buckets into a list of (height, width) tuples."""

if not buckets_str:

raise ValueError("Bucket string cannot be empty.")

bucket_pairs = buckets_str.strip().split(";")

parsed_buckets = []

for pair_str in bucket_pairs:

match = re.match(r"^\s*(\d+)\s*,\s*(\d+)\s*$", pair_str)

if not match:

raise ValueError(f"Invalid bucket format: '{pair_str}'. Expected 'height,width'.")

try:

height = int(match.group(1))

width = int(match.group(2))

if height <= 0 or width <= 0:

raise ValueError("Bucket dimensions must be positive integers.")

if height % 8 != 0 or width % 8 != 0:

warnings.warn(f"Bucket dimension ({height},{width}) not divisible by 8. This might cause issues.")

parsed_buckets.append((height, width))

except ValueError as e:

raise ValueError(f"Invalid integer in bucket pair '{pair_str}': {e}") from e

if not parsed_buckets:

raise ValueError("No valid buckets found in the provided string.")

"""Finds the closes bucket to the given height and width."""

min_metric = float("inf")

best_bucket_idx = None

for bucket_idx, (bucket_h, bucket_w) in enumerate(bucket_options):

metric = abs(h * bucket_w - w * bucket_h)

if metric <= min_metric:

min_metric = metric

best_bucket_idx = bucket_idx

(1, 1),

(4, 5),

(5, 4),

(2, 3),

(3, 2),

(3, 4),

(4, 3),

(9, 16),

(16, 9),

(1, 2),

(2, 1),

"""Generate a discrete list of ``(height, width)`` buckets for a target pixel budget.

target_pixels = resolution * resolution

if base_resolutions is not None:

aspect_ratios = [(w, h) for (h, w) in base_resolutions]

else:

aspect_ratios = _DEFAULT_BUCKET_ASPECT_RATIOS

buckets = []

seen = set()

for rw, rh in aspect_ratios:

aspect = rw / rh

# h * w == target_pixels and w / h == aspect => h = sqrt(target_pixels / aspect)

h = (target_pixels / aspect) ** 0.5

w = h * aspect

h = max(divisibility, round(h / divisibility) * divisibility)

w = max(divisibility, round(w / divisibility) * divisibility)

if (h, w) not in seen:

seen.add((h, w))

buckets.append((h, w))

"""

kwargs = {}

fsdp_state = getattr(accelerator.state, "fsdp_plugin", None)

if fsdp_state is None:

raise ValueError("Accelerate isn't configured to handle FSDP. Please update your installation.")

fsdp_plugin = accelerator.state.fsdp_plugin

if fsdp_plugin is None:

# FSDP not enabled in Accelerator

kwargs["sharding_strategy"] = ShardingStrategy.FULL_SHARD

else:

# FSDP is enabled → use plugin's strategy, or default if None

kwargs["sharding_strategy"] = fsdp_plugin.sharding_strategy or ShardingStrategy.FULL_SHARD

model: torch.nn.Module,

device: str | torch.device,

offload: bool = True,

use_orig_params: bool = True,

limit_all_gathers: bool = True,

fsdp_kwargs: dict[str, Any] | None = None,

transformer_layer_cls: set[type[torch.nn.Module]] | None = None,

) -> FSDP:

"""

logger = get_logger(__name__)

if transformer_layer_cls is None:

# Set the default layers if transformer_layer_cls is not provided

transformer_layer_cls = type(model.model.language_model.layers[0])

logger.info(f"transformer_layer_cls is not provided, auto-inferred as {transformer_layer_cls.__name__}")

# Add auto-wrap policy if transformer layers specified

auto_wrap_policy = partial(transformer_auto_wrap_policy, transformer_layer_cls={transformer_layer_cls})

config = {

"device_id": device,

"cpu_offload": CPUOffload(offload_params=offload) if offload else None,

"use_orig_params": use_orig_params,

"limit_all_gathers": limit_all_gathers,

"auto_wrap_policy": auto_wrap_policy,

}

if fsdp_kwargs:

config.update(fsdp_kwargs)

fsdp_model = FSDP(model, **config)

"""

def __init__(

self,

parameters: Iterable[torch.nn.Parameter],

decay: float = 0.9999,

min_decay: float = 0.0,

update_after_step: int = 0,

use_ema_warmup: bool = False,

inv_gamma: float | int = 1.0,

power: float | int = 2 / 3,

foreach: bool = False,

model_cls: Any | None = None,

model_config: dict[str, Any] | None = None,

**kwargs,

):

"""

if isinstance(parameters, torch.nn.Module):

deprecation_message = (

"Passing a `torch.nn.Module` to `ExponentialMovingAverage` is deprecated. "

"Please pass the parameters of the module instead."

)

deprecate(

"passing a `torch.nn.Module` to `ExponentialMovingAverage`",

"1.0.0",

deprecation_message,

standard_warn=False,

)

parameters = parameters.parameters()

# set use_ema_warmup to True if a torch.nn.Module is passed for backwards compatibility

use_ema_warmup = True

if kwargs.get("max_value", None) is not None:

deprecation_message = "The `max_value` argument is deprecated. Please use `decay` instead."

deprecate("max_value", "1.0.0", deprecation_message, standard_warn=False)

decay = kwargs["max_value"]

if kwargs.get("min_value", None) is not None:

deprecation_message = "The `min_value` argument is deprecated. Please use `min_decay` instead."

deprecate("min_value", "1.0.0", deprecation_message, standard_warn=False)

min_decay = kwargs["min_value"]

parameters = list(parameters)

self.shadow_params = [p.clone().detach() for p in parameters]

if kwargs.get("device", None) is not None:

deprecation_message = "The `device` argument is deprecated. Please use `to` instead."

deprecate("device", "1.0.0", deprecation_message, standard_warn=False)

self.to(device=kwargs["device"])

self.temp_stored_params = None

self.decay = decay

self.min_decay = min_decay

self.update_after_step = update_after_step

self.use_ema_warmup = use_ema_warmup

self.inv_gamma = inv_gamma

self.power = power

self.optimization_step = 0

self.cur_decay_value = None # set in `step()`

self.foreach = foreach

self.model_cls = model_cls

self.model_config = model_config

@classmethod

def from_pretrained(cls, path, model_cls, foreach=False) -> "EMAModel":

_, ema_kwargs = model_cls.from_config(path, return_unused_kwargs=True)

model = model_cls.from_pretrained(path)

ema_model = cls(model.parameters(), model_cls=model_cls, model_config=model.config, foreach=foreach)

ema_model.load_state_dict(ema_kwargs)

return ema_model

def save_pretrained(self, path):

if self.model_cls is None:

raise ValueError("`save_pretrained` can only be used if `model_cls` was defined at __init__.")

if self.model_config is None:

raise ValueError("`save_pretrained` can only be used if `model_config` was defined at __init__.")

model = self.model_cls.from_config(self.model_config)

state_dict = self.state_dict()

state_dict.pop("shadow_params", None)

model.register_to_config(**state_dict)

self.copy_to(model.parameters())

model.save_pretrained(path)

def get_decay(self, optimization_step: int) -> float:

"""

step = max(0, optimization_step - self.update_after_step - 1)

if step <= 0:

return 0.0

if self.use_ema_warmup:

cur_decay_value = 1 - (1 + step / self.inv_gamma) ** -self.power

else:

cur_decay_value = (1 + step) / (10 + step)

cur_decay_value = min(cur_decay_value, self.decay)

# make sure decay is not smaller than min_decay

cur_decay_value = max(cur_decay_value, self.min_decay)

return cur_decay_value

@torch.no_grad()

def step(self, parameters: Iterable[torch.nn.Parameter]):

if isinstance(parameters, torch.nn.Module):

deprecation_message = (

"Passing a `torch.nn.Module` to `ExponentialMovingAverage.step` is deprecated. "

"Please pass the parameters of the module instead."

)

deprecate(

"passing a `torch.nn.Module` to `ExponentialMovingAverage.step`",

"1.0.0",

deprecation_message,

standard_warn=False,

)

parameters = parameters.parameters()

parameters = list(parameters)

self.optimization_step += 1

# Compute the decay factor for the exponential moving average.

decay = self.get_decay(self.optimization_step)

self.cur_decay_value = decay

one_minus_decay = 1 - decay

context_manager = contextlib.nullcontext()

if self.foreach:

if is_transformers_available() and transformers.integrations.deepspeed.is_deepspeed_zero3_enabled():

context_manager = deepspeed.zero.GatheredParameters(parameters, modifier_rank=None)

with context_manager:

params_grad = [param for param in parameters if param.requires_grad]

s_params_grad = [

s_param for s_param, param in zip(self.shadow_params, parameters) if param.requires_grad

]

if len(params_grad) < len(parameters):

torch._foreach_copy_(

[s_param for s_param, param in zip(self.shadow_params, parameters) if not param.requires_grad],

[param for param in parameters if not param.requires_grad],

non_blocking=True,

)

torch._foreach_sub_(

s_params_grad, torch._foreach_sub(s_params_grad, params_grad), alpha=one_minus_decay

)

else:

for s_param, param in zip(self.shadow_params, parameters):

if is_transformers_available() and transformers.integrations.deepspeed.is_deepspeed_zero3_enabled():

context_manager = deepspeed.zero.GatheredParameters(param, modifier_rank=None)

with context_manager:

if param.requires_grad:

s_param.sub_(one_minus_decay * (s_param - param))

else:

s_param.copy_(param)

def copy_to(self, parameters: Iterable[torch.nn.Parameter]) -> None:

"""

parameters = list(parameters)

if self.foreach:

torch._foreach_copy_(

[param.data for param in parameters],

[s_param.to(param.device).data for s_param, param in zip(self.shadow_params, parameters)],

)

else:

for s_param, param in zip(self.shadow_params, parameters):

param.data.copy_(s_param.to(param.device).data)

def pin_memory(self) -> None:

r"""

self.shadow_params = [p.pin_memory() for p in self.shadow_params]

def to(self, device=None, dtype=None, non_blocking=False) -> None:

r"""

# .to() on the tensors handles None correctly

self.shadow_params = [

p.to(device=device, dtype=dtype, non_blocking=non_blocking)

if p.is_floating_point()

else p.to(device=device, non_blocking=non_blocking)

for p in self.shadow_params

]

def state_dict(self) -> dict:

r"""

# Following PyTorch conventions, references to tensors are returned:

# "returns a reference to the state and not its copy!" -

# https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict

return {

"decay": self.decay,

"min_decay": self.min_decay,

"optimization_step": self.optimization_step,

"update_after_step": self.update_after_step,

"use_ema_warmup": self.use_ema_warmup,

"inv_gamma": self.inv_gamma,

"power": self.power,

"shadow_params": self.shadow_params,

}

def store(self, parameters: Iterable[torch.nn.Parameter]) -> None:

r"""

self.temp_stored_params = [param.detach().cpu().clone() for param in parameters]

def restore(self, parameters: Iterable[torch.nn.Parameter]) -> None:

r"""

if self.temp_stored_params is None:

raise RuntimeError("This ExponentialMovingAverage has no `store()`ed weights to `restore()`")

if self.foreach:

torch._foreach_copy_(

[param.data for param in parameters], [c_param.data for c_param in self.temp_stored_params]

)

else:

for c_param, param in zip(self.temp_stored_params, parameters):

param.data.copy_(c_param.data)

# Better memory-wise.

self.temp_stored_params = None

def load_state_dict(self, state_dict: dict) -> None:

r"""

# deepcopy, to be consistent with module API

state_dict = copy.deepcopy(state_dict)

self.decay = state_dict.get("decay", self.decay)

if self.decay < 0.0 or self.decay > 1.0:

raise ValueError("Decay must be between 0 and 1")

self.min_decay = state_dict.get("min_decay", self.min_decay)

if not isinstance(self.min_decay, float):

raise ValueError("Invalid min_decay")

self.optimization_step = state_dict.get("optimization_step", self.optimization_step)

if not isinstance(self.optimization_step, int):

raise ValueError("Invalid optimization_step")

self.update_after_step = state_dict.get("update_after_step", self.update_after_step)

if not isinstance(self.update_after_step, int):

raise ValueError("Invalid update_after_step")

self.use_ema_warmup = state_dict.get("use_ema_warmup", self.use_ema_warmup)

if not isinstance(self.use_ema_warmup, bool):

raise ValueError("Invalid use_ema_warmup")

self.inv_gamma = state_dict.get("inv_gamma", self.inv_gamma)

if not isinstance(self.inv_gamma, (float, int)):

raise ValueError("Invalid inv_gamma")

self.power = state_dict.get("power", self.power)

if not isinstance(self.power, (float, int)):

raise ValueError("Invalid power")

shadow_params = state_dict.get("shadow_params", None)

if shadow_params is not None:

self.shadow_params = shadow_params

if not isinstance(self.shadow_params, list):

raise ValueError("shadow_params must be a list")

if not all(isinstance(p, torch.Tensor) for p in self.shadow_params):