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

rename ModelOnDevice to SwapToDevice

Browse Source
This commit is contained in:
Stephan Auerhahn
2023-08-06 23:46:20 +00:00
parent ced97f0e84
commit 6c18c8443a
1 changed files with 10 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
sgm/inference/helpers.py
View File

@@ -152,7 +152,7 @@ def do_sample(
with autocast(device) as precision_scope:
with model.ema_scope():
num_samples = [num_samples]
with ModelOnDevice(model.conditioner, device):
with SwapToDevice(model.conditioner, device):
batch, batch_uc = get_batch(
get_unique_embedder_keys_from_conditioner(model.conditioner),
value_dict,
@@ -189,11 +189,11 @@ def do_sample(
model.model, input, sigma, c, **additional_model_inputs
)
with ModelOnDevice(model.denoiser, device):
with ModelOnDevice(model.model, device):
with SwapToDevice(model.denoiser, device):
with SwapToDevice(model.model, device):
samples_z = sampler(denoiser, randn, cond=c, uc=uc)
with ModelOnDevice(model.first_stage_model, device):
with SwapToDevice(model.first_stage_model, device):
samples_x = model.decode_first_stage(samples_z)
samples = torch.clamp((samples_x + 1.0) / 2.0, min=0.0, max=1.0)
@@ -293,7 +293,7 @@ def do_img2img(
with torch.no_grad():
with autocast(device):
with model.ema_scope():
with ModelOnDevice(model.conditioner, device):
with SwapToDevice(model.conditioner, device):
batch, batch_uc = get_batch(
get_unique_embedder_keys_from_conditioner(model.conditioner),
value_dict,
@@ -313,7 +313,7 @@ def do_img2img(
if skip_encode:
z = img
else:
with ModelOnDevice(model.first_stage_model, device):
with SwapToDevice(model.first_stage_model, device):
z = model.encode_first_stage(img)
noise = torch.randn_like(z)
@@ -336,11 +336,11 @@ def do_img2img(
def denoiser(x, sigma, c):
return model.denoiser(model.model, x, sigma, c)
with ModelOnDevice(model.denoiser, device):
with ModelOnDevice(model.model, device):
with SwapToDevice(model.denoiser, device):
with SwapToDevice(model.model, device):
samples_z = sampler(denoiser, noised_z, cond=c, uc=uc)
with ModelOnDevice(model.first_stage_model, device):
with SwapToDevice(model.first_stage_model, device):
samples_x = model.decode_first_stage(samples_z)
samples = torch.clamp((samples_x + 1.0) / 2.0, min=0.0, max=1.0)
@@ -352,7 +352,7 @@ def do_img2img(
return samples
class ModelOnDevice(object):
class SwapToDevice(object):
def __init__(
self,
model: Union[torch.nn.Module, torch.Tensor],