manbeast3b commited on
Commit
f39847f
·
verified ·
1 Parent(s): b646e3a

Update src/pipeline.py

Browse files
Files changed (1) hide show
  1. src/pipeline.py +1 -1
src/pipeline.py CHANGED
@@ -56,7 +56,7 @@ def load_pipeline() -> Pipeline:
56
  pipeline = DiffusionPipeline.from_pretrained(ckpt_id, revision=ckpt_revision, transformer=transformer, text_encoder_2=text_encoder_2, torch_dtype=torch.bfloat16,)
57
  pipeline.to("cuda")
58
  pipeline.to(memory_format=torch.channels_last)
59
- for _ in range(1):
60
  pipeline(prompt="insensible, timbale, pothery, electrovital, actinogram, taxis, intracerebellar, centrodesmus", width=1024, height=1024, guidance_scale=0.0, num_inference_steps=4, max_sequence_length=256)
61
  return pipeline
62
 
 
56
  pipeline = DiffusionPipeline.from_pretrained(ckpt_id, revision=ckpt_revision, transformer=transformer, text_encoder_2=text_encoder_2, torch_dtype=torch.bfloat16,)
57
  pipeline.to("cuda")
58
  pipeline.to(memory_format=torch.channels_last)
59
+ with torch.inference_mode():
60
  pipeline(prompt="insensible, timbale, pothery, electrovital, actinogram, taxis, intracerebellar, centrodesmus", width=1024, height=1024, guidance_scale=0.0, num_inference_steps=4, max_sequence_length=256)
61
  return pipeline
62