Foreign-Mirrors/Fooocus
1
0
Fork 0
mirror of https://github.com/lllyasviel/Fooocus.git synced 2026-08-16 13:13:16 +02:00
Code Issues Packages Projects Releases Wiki Activity

try fix colab with virtual ram (#378)

Browse Source 
try fix colab with virtual ram (#378)
This commit is contained in:
lllyasviel
2023-09-15 01:24:07 -07:00
committed by GitHub
parent b5b4fd27f1
commit cf7cde08b1
7 changed files with 248 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/core.py
+11 -2
View File
@@ -10,6 +10,7 @@ import comfy.utils
from comfy.sd import load_checkpoint_guess_config
from nodes import VAEDecode, EmptyLatentImage
from comfy.sample import prepare_mask, broadcast_cond, load_additional_models, cleanup_additional_models
from comfy.model_base import SDXLRefiner
from modules.samplers_advanced import KSampler, KSamplerWithRefiner
from modules.patch import patch_all
@@ -20,7 +21,15 @@ opVAEDecode = VAEDecode()
class StableDiffusionModel:
def __init__(self, unet, vae, clip, clip_vision):
def __init__(self, unet, vae, clip, clip_vision, model_filename=None):
if isinstance(model_filename, str):
is_refiner = isinstance(unet.model, SDXLRefiner)
if unet is not None:
unet.model.model_file = dict(filename=model_filename, prefix='model')
if clip is not None:
clip.cond_stage_model.model_file = dict(filename=model_filename, prefix='refiner_clip' if is_refiner else 'base_clip')
if vae is not None:
vae.first_stage_model.model_file = dict(filename=model_filename, prefix='first_stage_model')
self.unet = unet
self.vae = vae
self.clip = clip
@@ -38,7 +47,7 @@ class StableDiffusionModel:
@torch.no_grad()
def load_model(ckpt_filename):
unet, clip, vae, clip_vision = load_checkpoint_guess_config(ckpt_filename)
return StableDiffusionModel(unet=unet, clip=clip, vae=vae, clip_vision=clip_vision)
return StableDiffusionModel(unet=unet, clip=clip, vae=vae, clip_vision=clip_vision, model_filename=ckpt_filename)
@torch.no_grad()