MolecularDiffusion.modules.models.loqi.self_conditioning

LoQI self-conditioning module (diffusion variant only).

Ported from others/LoQI/src/megalodon/models/self_conditioning.py (NVIDIA, Apache-2.0), narrowed to the single variable LoQI self-conditions on: x, with vector: True. In that mode the fusion is a 2-in / 1-out linear stack applied along a stacked last axis – i.e. a learned mix of the current noisy coordinates and the previous step’s prediction. loqi_flow.yaml has no self_conditioning block at all, so the flow task builds this as None.

modules_dict and its x key are load-bearing: they are the self_conditioning_module.*.modules_dict.x.* keys of the released weights.

Classes

BaseSelfConditioningModule

Fuse <var>_t with the previous step's <var>_hat.

Module Contents

class MolecularDiffusion.modules.models.loqi.self_conditioning.BaseSelfConditioningModule(variables: list[dict])

Bases: torch.nn.Module

Fuse <var>_t with the previous step’s <var>_hat.

forward(batch: dict, cond_batch: dict)

clamps: list[tuple] = []

fuse_softmax: list[bool] = []

keys: list[str] = []

modules_dict

vector_mask: list[bool] = []