fix micro cond gen in training

huggingface · Aug 10, 2023 · 85af31f · 85af31f
1 parent 46026a3
commit 85af31f
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/training/train_muse.py b/training/train_muse.py
@@ -940,8 +940,9 @@ def generate_images(
         clip_embeds = None
 
     if config.model.transformer.get("add_micro_cond_embeds", False):
+        resolution = config.dataset.preprocessing.resolution
         micro_conds = torch.tensor(
-            [256, 256, 0, 0, 6], device=encoder_hidden_states.device, dtype=encoder_hidden_states.dtype
+            [resolution, resolution, 0, 0, 6], device=encoder_hidden_states.device, dtype=encoder_hidden_states.dtype
         )
         micro_conds = micro_conds.unsqueeze(0).repeat(encoder_hidden_states.shape[0], 1)