Spaces:
Running on Zero
Running on Zero
| """Deterministic, RNG-isolated caption dropout for diffusion training.""" | |
| from __future__ import annotations | |
| import math | |
| from collections.abc import Mapping | |
| from numbers import Real | |
| import torch | |
| _CAPTION_DROPOUT_STREAM_ID = 0x4341_5054_494F_4E01 | |
| def validate_caption_dropout_probability(value: object) -> float: | |
| """Return a validated caption-dropout probability in ``[0, 1]``.""" | |
| if isinstance(value, bool) or not isinstance(value, Real): | |
| raise TypeError("caption_dropout_p must be a non-boolean real number.") | |
| probability = float(value) | |
| if not math.isfinite(probability) or not 0.0 <= probability <= 1.0: | |
| raise ValueError("caption_dropout_p must be finite and within [0, 1].") | |
| return probability | |
| def derive_caption_dropout_seed( | |
| base_seed: int, | |
| global_rank: int, | |
| global_step: int, | |
| ) -> int: | |
| """Derive a stable, position-addressed seed without consuming global RNG. | |
| The address deliberately includes rank and global training step. Batch-slot | |
| identity is supplied by the stable position within the mask vector sampled | |
| from this seed. | |
| """ | |
| components = (base_seed, global_rank, global_step, _CAPTION_DROPOUT_STREAM_ID) | |
| if any(isinstance(value, bool) or not isinstance(value, int) for value in components): | |
| raise TypeError("Caption-dropout seed components must be non-boolean integers.") | |
| if global_rank < 0 or global_step < 0: | |
| raise ValueError("Caption-dropout seed address requires global_rank and global_step >= 0.") | |
| mask = (1 << 64) - 1 | |
| value = 0x510E_527F_ADE6_82D1 | |
| for component in components: | |
| value ^= (component & mask) + 0x9E37_79B9_7F4A_7C15 + ((value << 6) & mask) | |
| value ^= value >> 30 | |
| value = (value * 0xBF58_476D_1CE4_E5B9) & mask | |
| value ^= value >> 27 | |
| value = (value * 0x94D0_49BB_1331_11EB) & mask | |
| value ^= value >> 31 | |
| return value & ((1 << 63) - 1) | |
| def sample_caption_dropout_mask( | |
| probability: object, | |
| batch_size: int, | |
| *, | |
| base_seed: int, | |
| global_rank: int, | |
| global_step: int, | |
| ) -> torch.Tensor: | |
| """Sample a per-sample CPU mask from an isolated generator. | |
| ``p=0`` and ``p=1`` use exact deterministic fast paths. For intermediate | |
| probabilities, a fresh CPU generator is addressed by seed/rank/step, so | |
| this function never reads or advances the process-global CPU or CUDA RNG. | |
| """ | |
| probability = validate_caption_dropout_probability(probability) | |
| if isinstance(batch_size, bool) or not isinstance(batch_size, int): | |
| raise TypeError("caption-dropout batch_size must be a non-boolean integer.") | |
| if batch_size <= 0: | |
| raise ValueError("caption-dropout batch_size must be positive.") | |
| seed = derive_caption_dropout_seed( | |
| base_seed=base_seed, | |
| global_rank=global_rank, | |
| global_step=global_step, | |
| ) | |
| if probability == 0.0: | |
| return torch.zeros(batch_size, dtype=torch.bool) | |
| if probability == 1.0: | |
| return torch.ones(batch_size, dtype=torch.bool) | |
| generator = torch.Generator(device="cpu") | |
| generator.manual_seed(seed) | |
| return torch.rand(batch_size, generator=generator) < probability | |
| def replace_prompt_embeds_with_negative( | |
| conditional_dict: Mapping[str, object], | |
| unconditional_dict: Mapping[str, object], | |
| dropout_mask: torch.Tensor, | |
| ) -> dict[str, object]: | |
| """Return conditioning with selected prompt-embedding rows replaced. | |
| Only ``prompt_embeds`` is changed. Other conditioning fields and both input | |
| mappings remain untouched. | |
| """ | |
| if not torch.is_tensor(dropout_mask): | |
| raise TypeError("caption-dropout mask must be a torch.Tensor.") | |
| if dropout_mask.dtype != torch.bool or dropout_mask.ndim != 1: | |
| raise ValueError("caption-dropout mask must be a one-dimensional bool tensor.") | |
| conditional = conditional_dict.get("prompt_embeds") | |
| negative = unconditional_dict.get("prompt_embeds") | |
| if not torch.is_tensor(conditional) or not torch.is_tensor(negative): | |
| raise KeyError("Both conditioning dictionaries must contain tensor prompt_embeds.") | |
| if conditional.shape != negative.shape: | |
| raise ValueError( | |
| "Conditional and negative prompt_embeds must have identical shapes, " | |
| f"got {tuple(conditional.shape)} and {tuple(negative.shape)}." | |
| ) | |
| if conditional.ndim < 1 or conditional.shape[0] != dropout_mask.numel(): | |
| raise ValueError( | |
| "caption-dropout mask length must equal the prompt_embeds batch dimension." | |
| ) | |
| if conditional.device != negative.device or conditional.dtype != negative.dtype: | |
| raise ValueError("Conditional and negative prompt_embeds must share device and dtype.") | |
| broadcast_shape = (dropout_mask.numel(),) + (1,) * (conditional.ndim - 1) | |
| selected = torch.where( | |
| dropout_mask.to(device=conditional.device).reshape(broadcast_shape), | |
| negative, | |
| conditional, | |
| ) | |
| result = dict(conditional_dict) | |
| result["prompt_embeds"] = selected | |
| return result | |