dnrops
/
ml-stable-diffusion
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Image encoder resources management fix (#182) 

* Unload image encoder for image2image in reduceMemory mode

* Include image encoder to resource management functions
This commit is contained in:
Oleg Ponomaryov 2023-05-07 07:22:19 +03:00 committed by GitHub
parent 940dba02ee
commit ea904e31d7
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 6 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
swift/StableDiffusion/pipeline/StableDiffusionPipeline.swift
View File

@ -101,6 +101,7 @@ public struct StableDiffusionPipeline: ResourceManaging {
try textEncoder.loadResources()
try unet.loadResources()
try decoder.loadResources()
try encoder?.loadResources()
try controlNet?.loadResources()
try safetyChecker?.loadResources()
}
@ -111,6 +112,7 @@ public struct StableDiffusionPipeline: ResourceManaging {
textEncoder.unloadResources()
unet.unloadResources()
decoder.unloadResources()
encoder?.unloadResources()
controlNet?.unloadResources()
safetyChecker?.unloadResources()
}
@ -120,6 +122,7 @@ public struct StableDiffusionPipeline: ResourceManaging {
try textEncoder.prewarmResources()
try unet.prewarmResources()
try decoder.prewarmResources()
try encoder?.prewarmResources()
try controlNet?.prewarmResources()
try safetyChecker?.prewarmResources()
}
@ -162,6 +165,9 @@ public struct StableDiffusionPipeline: ResourceManaging {
// Generate random latent samples from specified seed
var latents: [MLShapedArray<Float32>] = try generateLatentSamples(configuration: config, scheduler: scheduler[0])
if reduceMemory {
encoder?.unloadResources()
}
let timestepStrength: Float? = config.mode == .imageToImage ? config.strength : nil
// Convert cgImage for ControlNet into MLShapedArray