Spaces:
Runtime error
Runtime error
File size: 876 Bytes
82ef366 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 |
import torch
from diffusers import DDIMScheduler
from models.stable_diffusion import CrossImageAttentionStableDiffusionPipeline
from models.unet_2d_condition import FreeUUNet2DConditionModel
def get_stable_diffusion_model() -> CrossImageAttentionStableDiffusionPipeline:
print("Loading Stable Diffusion model...")
device = torch.device(f'cuda') if torch.cuda.is_available() else torch.device('cpu')
pipe = CrossImageAttentionStableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5",
safety_checker=None).to(device)
pipe.unet = FreeUUNet2DConditionModel.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="unet").to(device)
pipe.scheduler = DDIMScheduler.from_config("runwayml/stable-diffusion-v1-5", subfolder="scheduler")
print("Done.")
return pipe
|