mirror of
https://github.com/AUTOMATIC1111/stable-diffusion-webui.git
synced 2025-08-04 03:10:21 +00:00
add XL support for live previews: approx and TAESD
This commit is contained in:
@@ -2,9 +2,9 @@ import os
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
from modules import devices, paths
|
||||
from modules import devices, paths, shared
|
||||
|
||||
sd_vae_approx_model = None
|
||||
sd_vae_approx_models = {}
|
||||
|
||||
|
||||
class VAEApprox(nn.Module):
|
||||
@@ -31,19 +31,34 @@ class VAEApprox(nn.Module):
|
||||
return x
|
||||
|
||||
|
||||
def download_model(model_path, model_url):
|
||||
if not os.path.exists(model_path):
|
||||
os.makedirs(os.path.dirname(model_path), exist_ok=True)
|
||||
|
||||
print(f'Downloading VAEApprox model to: {model_path}')
|
||||
torch.hub.download_url_to_file(model_url, model_path)
|
||||
|
||||
|
||||
def model():
|
||||
global sd_vae_approx_model
|
||||
model_name = "vaeapprox-sdxl.pt" if getattr(shared.sd_model, 'is_sdxl', False) else "model.pt"
|
||||
loaded_model = sd_vae_approx_models.get(model_name)
|
||||
|
||||
if sd_vae_approx_model is None:
|
||||
model_path = os.path.join(paths.models_path, "VAE-approx", "model.pt")
|
||||
sd_vae_approx_model = VAEApprox()
|
||||
if loaded_model is None:
|
||||
model_path = os.path.join(paths.models_path, "VAE-approx", model_name)
|
||||
if not os.path.exists(model_path):
|
||||
model_path = os.path.join(paths.script_path, "models", "VAE-approx", "model.pt")
|
||||
sd_vae_approx_model.load_state_dict(torch.load(model_path, map_location='cpu' if devices.device.type != 'cuda' else None))
|
||||
sd_vae_approx_model.eval()
|
||||
sd_vae_approx_model.to(devices.device, devices.dtype)
|
||||
model_path = os.path.join(paths.script_path, "models", "VAE-approx", model_name)
|
||||
|
||||
return sd_vae_approx_model
|
||||
if not os.path.exists(model_path):
|
||||
model_path = os.path.join(paths.models_path, "VAE-approx", model_name)
|
||||
download_model(model_path, 'https://github.com/AUTOMATIC1111/stable-diffusion-webui/releases/download/v1.0.0-pre/' + model_name)
|
||||
|
||||
loaded_model = VAEApprox()
|
||||
loaded_model.load_state_dict(torch.load(model_path, map_location='cpu' if devices.device.type != 'cuda' else None))
|
||||
loaded_model.eval()
|
||||
loaded_model.to(devices.device, devices.dtype)
|
||||
sd_vae_approx_models[model_name] = loaded_model
|
||||
|
||||
return loaded_model
|
||||
|
||||
|
||||
def cheap_approximation(sample):
|
||||
|
Reference in New Issue
Block a user