flip.nvflare.controllers.scatter_and_gather_ldm

Classes

ScatterAndGatherLDM

Two-phase (autoencoder → diffusion model) FedAvg controller for latent diffusion.

Module Contents

class flip.nvflare.controllers.scatter_and_gather_ldm.ScatterAndGatherLDM(model_id: str = '', min_clients: int = 1, num_rounds_ae: int = 5, num_rounds_dm: int = 5, start_round: int = 0, model_locator_id: str = '', wait_time_after_min_received: int = 10, aggregator_id=AppConstants.DEFAULT_AGGREGATOR_ID, persistor_id=AppConstants.DEFAULT_PERSISTOR_ID, shareable_generator_id=AppConstants.DEFAULT_SHAREABLE_GENERATOR_ID, train_task_name=AppConstants.TASK_TRAIN, train_timeout: int = 0, ignore_result_error: bool = True, fatal_error_delay: int = 5, task_check_period: float = 0.5, persist_every_n_rounds: int = 1, memory_gc_rounds: int = 1)

Bases: flip.nvflare.controllers.scatter_and_gather.ScatterAndGather

Two-phase (autoencoder → diffusion model) FedAvg controller for latent diffusion.

Subclasses FLIP’s ScatterAndGather (itself a thin subclass of stock NVFLARE ScatterAndGather). Each instance runs a single phase — the job wires one controller per phase (train_task_name train_ae then train_dm) and NVFLARE runs them sequentially — so the round loop itself is stock’s, inherited verbatim (with its per-round aggregator.reset, failed-client tracking, memory_gc_rounds allocator cleanup, and any future upstream fixes), alongside all the shared FLIP hooks (_accept_train_result, the metrics relay, FlipEvents.ABORTED). Only what is genuinely LDM-specific is overridden:

• __init__() — two per-phase round counts (num_rounds_ae / num_rounds_dm) and a server-side model_locator for the phase-1 → phase-2 weight handoff. Delegates the shared FedAvg wiring/validation to the base, passing num_rounds = ae + dm so the inherited persist/state code sees a valid total before a phase starts.

• start_controller() — resolves the model_locator and applies the app config.json per-phase round-count overrides, then delegates component validation and initial-model loading to stock.

• control_flow() — points self._num_rounds at this phase’s round count (stock’s loop, NUM_ROUNDS prop/header and last-round persist all key off it); for train_dm first loads the phase-1 autoencoder weights via the persistor/model-locator; then delegates the round loop to stock.

• stop_controller() — cancels outstanding tasks (stock only flips the phase).

_num_rounds_ae = 5

_num_rounds_dm = 5

_fatal_error_delay = 5

model_locator_id = ''

model_locator = None

start_controller(fl_ctx: nvflare.apis.fl_context.FLContext) → None

control_flow(abort_signal: nvflare.apis.signal.Signal, fl_ctx: nvflare.apis.fl_context.FLContext) → None

stop_controller(fl_ctx: nvflare.apis.fl_context.FLContext) → None