skygpu
/
skynet
mirror of https://github.com/skygpu/skynet.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Only check if should cancel inference every two steps, also pipe to cuda if cpu offloading is off

pull/26/head
Guillermo Rodriguez 2023-10-05 23:20:45 -03:00
parent 9fa5a01c34
commit 93203ab533
No known key found for this signature in database
GPG Key ID: EC3AB66D5D83B392
2 changed files with 5 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
skynet/dgpu/compute.py
View File

@ -151,7 +151,7 @@ class SkynetMM:
num_inference_steps=step, num_inference_steps=step,
generator=seed, generator=seed,
callback=callback_fn, callback=callback_fn,
callback_steps=1, callback_steps=2,
**extra_params **extra_params
).images[0] ).images[0]

5
skynet/utils.py
View File

@ -113,6 +113,8 @@ def pipeline_for(model: str, mem_fraction: float = 1.0, image=False) -> Diffusio
pipe.scheduler = EulerAncestralDiscreteScheduler.from_config( pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(
pipe.scheduler.config) pipe.scheduler.config)
pipe.enable_xformers_memory_efficient_attention()
if over_mem: if over_mem:
if not image: if not image:
pipe.enable_vae_slicing() pipe.enable_vae_slicing()
@ -120,7 +122,8 @@ def pipeline_for(model: str, mem_fraction: float = 1.0, image=False) -> Diffusio
pipe.enable_model_cpu_offload() pipe.enable_model_cpu_offload()
pipe.enable_xformers_memory_efficient_attention() else:
pipe = pipe.to('cuda')
return pipe return pipe