Coverage for widitapp/apps.py: 100.00%
71 statements
« prev ^ index » next coverage.py v7.10.6, created at 2026-09-22 07:22 +0000
1from typing import Sequence
2from pathlib import Path
3from cluey import Cluey, main, tool, method
6class WiDiTApp(Cluey):
7 @method
8 def model(
9 self,
10 checkpoint: Path = None,
11 dim: int = 2,
12 unet: bool = False,
13 in_channels: int = 1,
14 out_channels: int = None,
15 use_diffusion: bool = True,
16 use_conditioning: bool = None,
17 verbose: bool = False,
18 preset: str = "",
19 hidden_size: int = 768,
20 depth: int = 12,
21 num_heads: int = 12,
22 patch_size: int = 2,
23 window_size: int = 4,
24 mlp_ratio: float = 4.0,
25 filters: int = 64,
26 kernel: int = 3,
27 layers: int = 4,
28 use_flash_attention: bool = True,
29 timestep_embed_dim: int = None,
30 **kwargs,
31 ):
32 if "spatial_dim" in kwargs:
33 dim = kwargs.pop("spatial_dim")
34 if "spatial_dims" in kwargs:
35 dim = kwargs.pop("spatial_dims")
37 if use_conditioning is None:
38 use_conditioning = use_diffusion
40 if out_channels is None:
41 out_channels = 1 + int(use_diffusion)
43 # If a checkpoint is provided, load the model from the checkpoint instead of building a new one
44 if checkpoint is not None:
45 from widit import load_model
47 model = load_model(checkpoint)
49 # TODO: Add some checks to make sure the loaded model is compatible with the provided parameters (e.g. in_channels, out_channels, etc.)
51 return model
53 if unet:
54 from widit import Unet
56 return Unet(
57 in_channels=in_channels,
58 out_channels=out_channels,
59 filters=filters,
60 kernel_size=kernel,
61 layers=layers,
62 spatial_dim=dim,
63 use_conditioning=use_conditioning,
64 timestep_embed_dim=timestep_embed_dim,
65 )
67 from widit import WiDiT, PRESETS
69 model_kwargs: dict = {}
71 if preset:
72 assert preset in PRESETS, f"Model '{preset}' not in PRESETS: {list(PRESETS.keys())}"
73 instantiator = PRESETS[preset]
74 else:
75 instantiator = WiDiT
76 model_kwargs = dict(
77 spatial_dim=dim,
78 hidden_size=hidden_size,
79 depth=depth,
80 num_heads=num_heads,
81 patch_size=patch_size,
82 )
84 model = instantiator(
85 in_channels=in_channels,
86 out_channels=out_channels,
87 use_conditioning=use_conditioning,
88 window_size=window_size,
89 mlp_ratio=mlp_ratio,
90 use_flash_attention=use_flash_attention,
91 timestep_embed_dim=timestep_embed_dim,
92 **model_kwargs
93 )
95 if verbose:
96 total = sum(p.numel() for p in model.parameters())
97 print(f"Model: {instantiator}")
98 print(f"Model Summary:\n{model}")
99 print(f"Model has {total:,} parameters")
100 trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
101 print(f"Model has {trainable:,} trainable parameters")
103 return model
105 @method
106 def datasets(
107 self,
108 **kwargs,
109 ) -> tuple:
110 """ Returns training and validation datasets """
111 raise NotImplementedError("Datasets method not yet implemented.")
113 @method("datasets")
114 def dataloaders(
115 self,
116 num_workers:int = 4,
117 batch_size:int = 1,
118 **kwargs,
119 ):
120 from torch.utils.data import DataLoader
122 training_dataset, validation_dataset = self.datasets(**kwargs)
123 training_dataloader = DataLoader(
124 training_dataset,
125 batch_size=batch_size,
126 shuffle=True,
127 drop_last=False,
128 num_workers=num_workers,
129 pin_memory=True,
130 )
132 validation_dataloader = None
133 if validation_dataset is not None:
134 validation_dataloader = DataLoader(
135 validation_dataset,
136 batch_size=batch_size,
137 shuffle=False,
138 drop_last=False,
139 num_workers=num_workers,
140 pin_memory=True,
141 )
143 return training_dataloader, validation_dataloader
145 @method
146 def loss(self, loss_fn: str = "mse", use_diffusion: bool = True):
147 """Build a supervised loss, or an optional extra diffusion image loss.
149 Override with ``@method`` to expose custom loss options to ``train``.
150 Return a module accepting (prediction, target). For diffusion these
151 are the unclipped clean-image estimate and clean target; the returned
152 loss is added to the standard diffusion objective. None adds nothing.
153 """
154 from .training import build_loss_fn
156 return None if use_diffusion else build_loss_fn(loss_fn)
158 @method
159 def metrics(self, use_diffusion: bool = True):
160 """Return named validation modules accepting (prediction, target).
162 Names omit the ``val/`` prefix. Each module returns a scalar batch
163 mean. Diffusion metrics use the unclipped clean-image estimate.
164 Extend with ``@method("super")`` and ``super().metrics(**kwargs)``.
165 """
166 import torch
168 if use_diffusion:
169 return {}
170 return {"mse": torch.nn.MSELoss(), "smoothl1": torch.nn.SmoothL1Loss()}
172 @tool("model", "dataloaders", "loss", "metrics")
173 def train(
174 self,
175 epochs: int = 40,
176 log_every: int = 100,
177 learning_rate: float = 1e-4,
178 results_dir: Path = "./results",
179 use_diffusion: bool = True,
180 wandb: bool = False,
181 wandb_project: str = "",
182 run_name: str = "",
183 **kwargs,
184 ):
185 """ Train the model """
187 from .training import train
189 model = self.model(use_diffusion=use_diffusion, **kwargs)
191 wandb_project = wandb_project or str(self.__class__.__name__)
193 training_dataloader, validation_dataloader = self.dataloaders(**kwargs)
194 criterion = self.loss(use_diffusion=use_diffusion, **kwargs)
195 loss_kwargs = (
196 {"diffusion_loss_fn": criterion}
197 if use_diffusion else {"loss_fn": criterion}
198 )
199 train(
200 model=model,
201 training_dataloader=training_dataloader,
202 validation_dataloader=validation_dataloader,
203 results_dir=results_dir,
204 use_diffusion=use_diffusion,
205 learning_rate=learning_rate,
206 epochs=epochs,
207 log_every=log_every,
208 run_name=run_name,
209 wandb_logging=wandb,
210 wandb_project=wandb_project,
211 metrics=self.metrics(use_diffusion=use_diffusion, **kwargs),
212 **loss_kwargs,
213 )
215 @main
216 def predict(
217 self,
218 **kwargs,
219 ):
220 """ Makes predictions """
221 raise NotImplementedError("Prediction not yet implemented.")