Update all files for DiffusionSat-Single-256

unet/sat_unet.py

@@ -18,8 +18,23 @@ class SatUNet(UNet2DConditionModel):
 
     _supports_gradient_checkpointing = True
 
-    def __init__(self, *args, use_metadata: bool = True, num_metadata: int = 7, **kwargs):
-        super().__init__(*args, **kwargs)
+    @classmethod
+    def from_config(cls, config: Optional[Dict[str, Any]] = None, **kwargs):
+        """Load config using parent's expected keys so all UNet params (e.g. cross_attention_dim, use_linear_projection) are passed through."""
+        if config is None:
+            raise ValueError("Please provide a config.")
+        if not isinstance(config, dict):
+            config, kwargs = cls.load_config(config, return_unused_kwargs=True, **kwargs)
+        # Use parent's extract_init_dict so all UNet2DConditionModel params are in init_dict
+        init_dict, unused_kwargs, hidden_dict = UNet2DConditionModel.extract_init_dict(config, **kwargs)
+        init_dict.setdefault("use_metadata", True)
+        init_dict.setdefault("num_metadata", 7)
+        model = cls(**init_dict)
+        model.register_to_config(**hidden_dict)
+        return model
+
+    def __init__(self, *args, use_metadata: bool = True, num_metadata: int = 7, use_linear_projection: bool = False, **kwargs):
+        super().__init__(*args, use_linear_projection=use_linear_projection, **kwargs)
 
         # Track custom config entries for save/load parity with the base model.
         self.register_to_config(use_metadata=use_metadata, num_metadata=num_metadata)