Skip to content

Instantly share code, notes, and snippets.

@mchaker

mchaker/improve-stablediffusion-pipeline-mem-mgmt.py

Forked from fladdict/gist:2115eb7ea32c9245e4f45642553aa3e9
Created September 6, 2022 09:58
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save mchaker/a8125bfb0ab74aabde0af4544453d402 to your computer and use it in GitHub Desktop.
Save mchaker/a8125bfb0ab74aabde0af4544453d402 to your computer and use it in GitHub Desktop.
Download ZIP
Improve Diffuser Pipeline / StableDiffusion Memory Management
Raw
improve-stablediffusion-pipeline-mem-mgmt.py
pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", revision="fp16", torch_dtype=torch.float16, use_auth_token=True)
pipe = pipe.to("cuda")
img2imgPipe = StableDiffusionImg2ImgPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", revision="fp16", torch_dtype=torch.float16, use_auth_token=True)
img2imgPipe = img2imgPipe.to("cuda")
inpaintingPipe = StableDiffusionInpaintingPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", revision="fp16", torch_dtype=torch.float16, use_auth_token=True)
inpaintingPipe = inpaintingPipe.to("cuda")
#Hack the instance to reduce memory
#Pipes can share models
img2imgPipe.vae = inpaintingPipe.vae = pipe.vae
img2imgPipe.text_encoder = inpaintingPipe.text_encoder = pipe.text_encoder
img2imgPipe.tokenizer = inpaintingPipe.tokenizer = pipe.tokenizer
img2imgPipe.safety_checker = inpaintingPipe.safety_checker = pipe.safety_checker
img2imgPipe.unet = inpaintingPipe.unet = pipe.unet
img2imgPipe.feature_extractor = inpaintingPipe.feature_extractor = pipe.feature_extractor
img2imgPipe.scheduler = inpaintingPipe.scheduler = pipe.scheduler
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment