(self)
| 451 | @require_transformers_version_greater("4.44.0") |
| 452 | class SlowBnb4BitTests(Base4bitTests): |
| 453 | def setUp(self) -> None: |
| 454 | gc.collect() |
| 455 | backend_empty_cache(torch_device) |
| 456 | |
| 457 | nf4_config = BitsAndBytesConfig( |
| 458 | load_in_4bit=True, |
| 459 | bnb_4bit_quant_type="nf4", |
| 460 | bnb_4bit_compute_dtype=torch.float16, |
| 461 | ) |
| 462 | model_4bit = SD3Transformer2DModel.from_pretrained( |
| 463 | self.model_name, subfolder="transformer", quantization_config=nf4_config, device_map=torch_device |
| 464 | ) |
| 465 | self.pipeline_4bit = DiffusionPipeline.from_pretrained( |
| 466 | self.model_name, transformer=model_4bit, torch_dtype=torch.float16 |
| 467 | ) |
| 468 | self.pipeline_4bit.enable_model_cpu_offload() |
| 469 | |
| 470 | def tearDown(self): |
| 471 | del self.pipeline_4bit |
nothing calls this directly
no test coverage detected