jokerbit commited on
Commit
ac14b3c
·
verified ·
1 Parent(s): a2814a9

CUDA first

Browse files
Files changed (1) hide show
  1. src/pipeline.py +1 -2
src/pipeline.py CHANGED
@@ -42,12 +42,11 @@ def load_pipeline() -> Pipeline:
42
  local_files_only=True,
43
  torch_dtype=torch.bfloat16,
44
  )
45
-
46
  pipeline.transformer.to(memory_format=torch.channels_last)
47
  pipeline.transformer = torch.compile(pipeline.transformer)
48
  # pipeline.text_encoder.fuse_qkv_projections()
49
  # pipeline.vae = torch.compile(pipeline.vae)
50
- pipeline.to("cuda")
51
  for _ in range(4):
52
  pipeline("cat", num_inference_steps=4)
53
  torch.cuda.empty_cache()
 
42
  local_files_only=True,
43
  torch_dtype=torch.bfloat16,
44
  )
45
+ pipeline.to("cuda")
46
  pipeline.transformer.to(memory_format=torch.channels_last)
47
  pipeline.transformer = torch.compile(pipeline.transformer)
48
  # pipeline.text_encoder.fuse_qkv_projections()
49
  # pipeline.vae = torch.compile(pipeline.vae)
 
50
  for _ in range(4):
51
  pipeline("cat", num_inference_steps=4)
52
  torch.cuda.empty_cache()