Skip to content

Instantly share code, notes, and snippets.

@fladdict

fladdict/gist:2115eb7ea32c9245e4f45642553aa3e9

Last active July 29, 2023 19:40
Show Gist options
  • Star 5 You must be signed in to star a gist
  • Fork 1 You must be signed in to fork a gist
  • Save fladdict/2115eb7ea32c9245e4f45642553aa3e9 to your computer and use it in GitHub Desktop.
Save fladdict/2115eb7ea32c9245e4f45642553aa3e9 to your computer and use it in GitHub Desktop.
Download ZIP
Improve Diffuser Pipeline / StableDiffusion Memory Management
Raw
gistfile1.txt
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
@parlance-zz
Copy link

Works for me; lowers memory consumption by about 20% and all 3 pipelines still work.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment