Reduce memory usage for sd 3.5. (#2582)

This commit is contained in:
Laurent Mazare 2024-10-28 22:45:02 +01:00 committed by GitHub
parent 498bc2cdc9
commit 139ff56aeb
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 2 additions and 0 deletions

View File

@ -180,6 +180,8 @@ fn main() -> Result<()> {
let (context, y) = triple.encode_text_to_embedding(prompt.as_str(), &device)?;
let (context_uncond, y_uncond) =
triple.encode_text_to_embedding(uncond_prompt.as_str(), &device)?;
// Drop the text model early to avoid using too much memory.
drop(triple);
let context = Tensor::cat(&[context, context_uncond], 0)?;
let y = Tensor::cat(&[y, y_uncond], 0)?;