fix: fix diffusion example [DET-10470] (#10161)

determined-ai · Oct 31, 2024 · 85a5d22 · 85a5d22
1 parent 06b8f48
commit 85a5d22
Show file tree

Hide file tree

Showing 3 changed files with 2 additions and 2 deletions.
diff --git a/e2e_tests/tests/nightly/test_distributed.py b/e2e_tests/tests/nightly/test_distributed.py
@@ -114,6 +114,7 @@ def test_textual_inversion_stable_diffusion_generate() -> None:
     config["hyperparameters"]["num_batches"] = 2
     prompt_vals = config["hyperparameters"]["call_kwargs"]["prompt"]["vals"]
     config["hyperparameters"]["call_kwargs"]["guidance_scale"] = 7.5
+    config["max_restarts"] = 0
     while len(prompt_vals) > 1:
         prompt_vals.pop()
 

diff --git a/examples/diffusion/textual_inversion_stable_diffusion/detsd/pipeline.py b/examples/diffusion/textual_inversion_stable_diffusion/detsd/pipeline.py
@@ -211,7 +211,7 @@ def generate_on_cluster(cls) -> None:
                             devices_and_generators=devices_and_generators,
                             trial_id=trial_id,
                         )
-                        core_context.train.report_progress(pipeline.steps_completed / MAX_LENGTH)
+                        core_context.train.report_progress(pipeline.steps_completed / num_batches)
 
                     # Only preempt after a checkpoint has been saved.
                     if core_context.preempt.should_preempt():

diff --git a/examples/diffusion/textual_inversion_stable_diffusion/generate_grid.yaml b/examples/diffusion/textual_inversion_stable_diffusion/generate_grid.yaml
@@ -5,7 +5,6 @@ entrypoint: python3 -m determined.launch.torch_distributed python3 generate.py
 searcher:
   name: grid
   metric: none
-  max_length: 5 # Number of times the generation pipeline is called, per-worker.
 resources:
   slots_per_trial: 2 # Also equal to the number of unique initial seeds which will be used.
   max_slots: 4