Add files using upload-large-folder tool

pipeline_diffusionsat.py

@@ -161,6 +161,7 @@ class DiffusionSatPipeline(DiffusionPipeline):
     disable_vae_slicing = DiffusersStableDiffusionPipeline.disable_vae_slicing
     enable_sequential_cpu_offload = DiffusersStableDiffusionPipeline.enable_sequential_cpu_offload
     _execution_device = DiffusersStableDiffusionPipeline._execution_device
+    encode_prompt = DiffusersStableDiffusionPipeline.encode_prompt
     _encode_prompt = DiffusersStableDiffusionPipeline._encode_prompt
     run_safety_checker = DiffusersStableDiffusionPipeline.run_safety_checker
     decode_latents = DiffusersStableDiffusionPipeline.decode_latents

unet/sat_unet.py

@@ -74,7 +74,9 @@ class SatUNet(UNet2DConditionModel):
         # Reuse the same projection used for timestep encoding to stay aligned with base embeddings.
         projected = self.time_proj(metadata.view(-1)).view(md_bsz, self.num_metadata, -1).to(dtype=dtype)
 
-        md_emb = projected.new_zeros((md_bsz, projected.shape[-1]))
+        # md_embed outputs time_embed_dim (1280), not projected.shape[-1] (320)
+        time_embed_dim = self.time_embedding.linear_2.out_features
+        md_emb = torch.zeros((md_bsz, time_embed_dim), device=metadata.device, dtype=dtype)
         for idx, md_embed in enumerate(self.metadata_embedding):
             md_emb = md_emb + md_embed(projected[:, idx, :])