English
John6666 commited on
Commit
9c63f51
·
verified ·
1 Parent(s): 8efcd5c

Upload 2 files

Browse files
Files changed (1) hide show
  1. handler.py +2 -2
handler.py CHANGED
@@ -152,14 +152,14 @@ def load_pipeline_opt(repo_id: str, dtype: torch.dtype) -> Any:
152
  elif IS_CC89: quantize_(transformer, float8_dynamic_activation_float8_weight(), device="cuda")
153
  else: quantize_(transformer, weight, device="cuda")
154
  transformer.to(memory_format=torch.channels_last)
155
- transformer = torch.compile(transformer, mode="max-autotune", fullgraph=True)
156
  vae = AutoencoderKL.from_pretrained(repo_id, subfolder="vae", torch_dtype=dtype)
157
  vae.fuse_qkv_projections()
158
  if IS_CC90: quantize_(vae, float8_dynamic_activation_float8_weight(granularity=PerRow()), device="cuda")
159
  elif IS_CC89: quantize_(vae, float8_dynamic_activation_float8_weight(), device="cuda")
160
  else: quantize_(vae, weight, device="cuda")
161
  vae.to(memory_format=torch.channels_last)
162
- vae = torch.compile(vae, mode="max-autotune", fullgraph=True)
163
  pipe = FluxPipeline.from_pretrained(repo_id, transformer=None, vae=None, torch_dtype=dtype, quantization_config=quantization_config)
164
  pipe.transformer = transformer
165
  pipe.vae = vae
 
152
  elif IS_CC89: quantize_(transformer, float8_dynamic_activation_float8_weight(), device="cuda")
153
  else: quantize_(transformer, weight, device="cuda")
154
  transformer.to(memory_format=torch.channels_last)
155
+ transformer = torch.compile(transformer, mode="reduce-overhead", fullgraph=True)
156
  vae = AutoencoderKL.from_pretrained(repo_id, subfolder="vae", torch_dtype=dtype)
157
  vae.fuse_qkv_projections()
158
  if IS_CC90: quantize_(vae, float8_dynamic_activation_float8_weight(granularity=PerRow()), device="cuda")
159
  elif IS_CC89: quantize_(vae, float8_dynamic_activation_float8_weight(), device="cuda")
160
  else: quantize_(vae, weight, device="cuda")
161
  vae.to(memory_format=torch.channels_last)
162
+ vae = torch.compile(vae, mode="reduce-overhead", fullgraph=True)
163
  pipe = FluxPipeline.from_pretrained(repo_id, transformer=None, vae=None, torch_dtype=dtype, quantization_config=quantization_config)
164
  pipe.transformer = transformer
165
  pipe.vae = vae