avoid deep copies from dc.asdict()

Author: catwell

src/finegrain_toolbox/dc.py

@@ -0,0 +1,15 @@
+import dataclasses as dc
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from _typeshed import DataclassInstance
+
+
+def shallow_asdict(m: "DataclassInstance") -> dict[str, Any]:
+    return {field.name: getattr(m, field.name) for field in dc.fields(m)}
+
+
+class DcMixin:
+    def shallow_asdict(self) -> dict[str, Any]:
+        assert dc.is_dataclass(self)
+        return shallow_asdict(self)

src/finegrain_toolbox/flux/model.py

@@ -7,6 +7,7 @@
 import torch
 from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler, FluxTransformer2DModel
 
+from ..dc import DcMixin
 from ..models import SafePushToHubMixin
 from ..torch import default_device, default_dtype
 from ..types import Self
@@ -25,7 +26,7 @@ def get_mu(scheduler: FlowMatchEulerDiscreteScheduler, image_seq_len: int) -> fl
 
 
 @dc.dataclass(kw_only=True)
-class Model(SafePushToHubMixin):
+class Model(SafePushToHubMixin, DcMixin):
     device: torch.device
     dtype: torch.dtype
     scheduler: FlowMatchEulerDiscreteScheduler
@@ -62,7 +63,7 @@ def to(
         device: torch.device | None = None,
         dtype: torch.dtype | None = None,
     ) -> Self:
-        params = dc.asdict(self)
+        params = self.shallow_asdict()
         if device is not None:
             params["device"] = device
         if dtype is not None:

src/finegrain_toolbox/flux/prompt.py

@@ -4,6 +4,7 @@
 import torch
 from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5TokenizerFast
 
+from ..dc import DcMixin
 from ..torch import default_device, default_dtype
 from ..types import Self
 
@@ -16,7 +17,7 @@ class Tokenized:
 
 
 @dc.dataclass(kw_only=True)
-class Prompt:
+class Prompt(DcMixin):
     t5_embeds: torch.Tensor  # (batch_size, seq_len, 4096)
     clip_embeds: torch.Tensor  # (batch_size, 768)
     text_ids: torch.Tensor  # shape (seq_len, 3)
@@ -38,7 +39,7 @@ def to(
         device: torch.device | None = None,
         dtype: torch.dtype | None = None,
     ) -> Self:
-        params = dc.asdict(self)
+        params = self.shallow_asdict()
         params["t5_embeds"] = self.t5_embeds.to(device=device, dtype=dtype)
         params["clip_embeds"] = self.clip_embeds.to(device=device, dtype=dtype)
         params["text_ids"] = self.text_ids.to(device=device, dtype=dtype)