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

path helper & model swapping rewrite

Browse Source
This commit is contained in:
Stephan Auerhahn
2023-08-10 04:35:59 -07:00
parent fc498bfaef
commit e190ecc60b
4 changed files with 94 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
scripts/demo/streamlit_helpers.py
View File

@@ -20,9 +20,7 @@ from sgm.inference.api import (
SamplingPipeline,
Thresholder,
)
from sgm.inference.helpers import (
embed_watermark,
)
from sgm.inference.helpers import embed_watermark, CudaModelLoader
@st.cache_resource()
@@ -35,10 +33,12 @@ def init_st(spec: SamplingSpec, load_ckpt=True, load_filter=True) -> Dict[str, A
if lowvram_mode:
pipeline = SamplingPipeline(
model_spec=spec, use_fp16=True, device="cuda", swap_device="cpu"
model_spec=spec,
use_fp16=True,
model_loader=CudaModelLoader(device="cuda", swap_device="cpu"),
)
else:
pipeline = SamplingPipeline(model_spec=spec, use_fp16=True, device="cuda")
pipeline = SamplingPipeline(model_spec=spec, use_fp16=False)
state["spec"] = spec
state["model"] = pipeline