Update src/pipeline.py
Browse files- 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 |
-
|
| 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 |
|