Method setup

src/diffusers/models/_modeling_parallel.py:174–187

(
        self,
        rank: int,
        world_size: int,
        device: torch.device,
        *,
        mesh: torch.distributed.device_mesh.DeviceMesh | None = None,
    )

    _mesh: torch.distributed.device_mesh.DeviceMesh = None

    def setup(
        self,
        rank: int,
        world_size: int,
        device: torch.device,
        *,
        mesh: torch.distributed.device_mesh.DeviceMesh | None = None,
    ):
        self._rank = rank
        self._world_size = world_size
        self._device = device
        self._mesh = mesh
        if self.context_parallel_config is not None:
            self.context_parallel_config.setup(rank, world_size, device, mesh)


@dataclass(frozen=True)