fix the precision bug in inference

NVlabs · Jan 5, 2025 · ac55992 · ac55992
1 parent db768da
commit ac55992
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 2 deletions.
diff --git a/configs/sana_config/512ms/Sana_600M_img512.yaml b/configs/sana_config/512ms/Sana_600M_img512.yaml
@@ -66,7 +66,7 @@ scheduler:
   predict_v: true
   noise_schedule: linear_flow
   pred_sigma: false
-  flow_shift: 1.0
+  flow_shift: 3.0
   # logit-normal timestep
   weighting_scheme: logit_normal
   logit_mean: 0.0

diff --git a/scripts/infer_run_inference_geneval.sh b/scripts/infer_run_inference_geneval.sh
@@ -72,7 +72,7 @@ echo "Add label: $add_label"
 echo "Exist time prefix: $exist_time_prefix"
 
 cmd_template="DPM_TQDM=True python scripts/inference_geneval.py --config={config_file} --model_path={model_path} \
-    --sampling_algo $sampling_algo --step $step --cfg_scale $cfg_scale \
+    --sampling_algo $sampling_algo --step $step --cfg_scale $cfg_scale --sample_nums $sample_nums \
     --gpu_id {gpu_id} --start_index {start_index} --end_index {end_index}"
 if [ -n "${add_label}" ]; then
     cmd_template="${cmd_template} --add_label ${add_label}"
@@ -108,6 +108,7 @@ if [[ "$model_paths" == *.pth ]]; then
     cmd="${cmd//\{end_index\}/$end_index}"
 
     echo "Running on GPU $gpu_id: samples $start_index to $end_index"
+    echo $cmd
     eval CUDA_VISIBLE_DEVICES=$gpu_id $cmd &
   done
   wait

diff --git a/scripts/inference_geneval.py b/scripts/inference_geneval.py
@@ -235,6 +235,7 @@ def visualize(sample_steps, cfg_scale, pag_scale):
                         latent_size,
                         device=device,
                         generator=generator,
+                        dtype=weight_dtype,
                     )
                     model_kwargs = dict(data_info={"img_hw": hw, "aspect_ratio": ar}, mask=emb_masks)