aljaz/generative-models
1
0
Fork 0
mirror of https://github.com/Stability-AI/generative-models.git synced 2026-02-02 04:14:27 +01:00
Code Issues Packages Projects Releases Wiki Activity

extract path resolution method, fix/improve device swapping support

Browse Source
This commit is contained in:
Stephan Auerhahn
2023-08-09 19:31:59 -07:00
parent d245e2002f
commit b51c36b0df
2 changed files with 45 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
scripts/demo/streamlit_helpers.py
View File

@@ -33,11 +33,12 @@ def init_st(spec: SamplingSpec, load_ckpt=True, load_filter=True) -> Dict[str, A
config = spec.config
ckpt = spec.ckpt
pipeline = SamplingPipeline(
model_spec=spec,
use_fp16=lowvram_mode,
device="cpu" if lowvram_mode else "cuda",
)
if lowvram_mode:
pipeline = SamplingPipeline(
model_spec=spec, use_fp16=True, device="cuda", swap_device="cpu"
)
else:
pipeline = SamplingPipeline(model_spec=spec, use_fp16=True, device="cuda")
state["spec"] = spec
state["model"] = pipeline