@@ -1210,11 +1210,11 @@ def get_latents(self, seed=0, shape=(4, 4, 64, 64), fp16=False):
1210
1210
return image
1211
1211
1212
1212
def get_unet_model (self , fp16 = False , model_id = "CompVis/stable-diffusion-v1-4" ):
1213
- revision = "fp16" if fp16 else None
1213
+ variant = "fp16" if fp16 else None
1214
1214
torch_dtype = torch .float16 if fp16 else torch .float32
1215
1215
1216
1216
model = UNet2DConditionModel .from_pretrained (
1217
- model_id , subfolder = "unet" , torch_dtype = torch_dtype , revision = revision
1217
+ model_id , subfolder = "unet" , torch_dtype = torch_dtype , variant = variant
1218
1218
)
1219
1219
model .to (torch_device ).eval ()
1220
1220
@@ -1376,7 +1376,7 @@ def test_compvis_sd_v1_4_fp16(self, seed, timestep, expected_slice):
1376
1376
@require_torch_accelerator
1377
1377
@skip_mps
1378
1378
def test_compvis_sd_v1_5 (self , seed , timestep , expected_slice ):
1379
- model = self .get_unet_model (model_id = "Jiali /stable-diffusion-1. 5" )
1379
+ model = self .get_unet_model (model_id = "stable-diffusion-v1-5 /stable-diffusion-v1- 5" )
1380
1380
latents = self .get_latents (seed )
1381
1381
encoder_hidden_states = self .get_encoder_hidden_states (seed )
1382
1382
@@ -1404,7 +1404,7 @@ def test_compvis_sd_v1_5(self, seed, timestep, expected_slice):
1404
1404
)
1405
1405
@require_torch_accelerator_with_fp16
1406
1406
def test_compvis_sd_v1_5_fp16 (self , seed , timestep , expected_slice ):
1407
- model = self .get_unet_model (model_id = "Jiali /stable-diffusion-1. 5" , fp16 = True )
1407
+ model = self .get_unet_model (model_id = "stable-diffusion-v1-5 /stable-diffusion-v1- 5" , fp16 = True )
1408
1408
latents = self .get_latents (seed , fp16 = True )
1409
1409
encoder_hidden_states = self .get_encoder_hidden_states (seed , fp16 = True )
1410
1410
@@ -1433,7 +1433,7 @@ def test_compvis_sd_v1_5_fp16(self, seed, timestep, expected_slice):
1433
1433
@require_torch_accelerator
1434
1434
@skip_mps
1435
1435
def test_compvis_sd_inpaint (self , seed , timestep , expected_slice ):
1436
- model = self .get_unet_model (model_id = "botp/ stable-diffusion-v1-5-inpainting" )
1436
+ model = self .get_unet_model (model_id = "stable-diffusion-v1-5/stable-diffusion -inpainting" )
1437
1437
latents = self .get_latents (seed , shape = (4 , 9 , 64 , 64 ))
1438
1438
encoder_hidden_states = self .get_encoder_hidden_states (seed )
1439
1439
@@ -1461,7 +1461,7 @@ def test_compvis_sd_inpaint(self, seed, timestep, expected_slice):
1461
1461
)
1462
1462
@require_torch_accelerator_with_fp16
1463
1463
def test_compvis_sd_inpaint_fp16 (self , seed , timestep , expected_slice ):
1464
- model = self .get_unet_model (model_id = "botp/ stable-diffusion-v1-5-inpainting" , fp16 = True )
1464
+ model = self .get_unet_model (model_id = "stable-diffusion-v1-5/stable-diffusion -inpainting" , fp16 = True )
1465
1465
latents = self .get_latents (seed , shape = (4 , 9 , 64 , 64 ), fp16 = True )
1466
1466
encoder_hidden_states = self .get_encoder_hidden_states (seed , fp16 = True )
1467
1467
0 commit comments