Skip to content

Commit

Permalink
Change tqdm disable conditions and deepcopy model map on init.
Browse files Browse the repository at this point in the history
  • Loading branch information
eagarvey-amd committed Jul 17, 2024
1 parent 2de912e commit 9fdc07f
Showing 1 changed file with 4 additions and 2 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -247,7 +247,7 @@ def __init__(
"external_weights": None,
"external_weight_path": None,
}
sd_model_map = get_sd_model_map(hf_model_name)
sd_model_map = copy.deepcopy(get_sd_model_map(hf_model_name))
for submodel in sd_model_map:
if "load" not in sd_model_map[submodel]:
sd_model_map[submodel]["load"] = True
Expand Down Expand Up @@ -569,9 +569,11 @@ def _produce_latents_sdxl(
[guidance_scale],
dtype=self.map["unet"]["np_dtype"],
)
# Disable progress bar if we aren't in verbose mode or if we're printing
# benchmark latencies for unet.
for i, t in tqdm(
enumerate(timesteps),
disable=(self.map["unet"].get("benchmark") and self.verbose),
disable=(self.map["unet"].get("benchmark") or not self.verbose),
):
if self.cpu_scheduling:
latent_model_input, t = self.scheduler.scale_model_input(
Expand Down

0 comments on commit 9fdc07f

Please sign in to comment.