mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2025-01-11 10:25:16 +00:00
Fix imports.
This commit is contained in:
parent
7e51bbd07f
commit
bae4fb4a9d
@ -5,17 +5,17 @@ import torch
|
|||||||
import torch as th
|
import torch as th
|
||||||
import torch.nn as nn
|
import torch.nn as nn
|
||||||
|
|
||||||
from ldm.modules.diffusionmodules.util import (
|
from ..ldm.modules.diffusionmodules.util import (
|
||||||
conv_nd,
|
conv_nd,
|
||||||
linear,
|
linear,
|
||||||
zero_module,
|
zero_module,
|
||||||
timestep_embedding,
|
timestep_embedding,
|
||||||
)
|
)
|
||||||
|
|
||||||
from ldm.modules.attention import SpatialTransformer
|
from ..ldm.modules.attention import SpatialTransformer
|
||||||
from ldm.modules.diffusionmodules.openaimodel import UNetModel, TimestepEmbedSequential, ResBlock, Downsample, AttentionBlock
|
from ..ldm.modules.diffusionmodules.openaimodel import UNetModel, TimestepEmbedSequential, ResBlock, Downsample, AttentionBlock
|
||||||
from ldm.models.diffusion.ddpm import LatentDiffusion
|
from ..ldm.models.diffusion.ddpm import LatentDiffusion
|
||||||
from ldm.util import log_txt_as_img, exists, instantiate_from_config
|
from ..ldm.util import log_txt_as_img, exists, instantiate_from_config
|
||||||
|
|
||||||
|
|
||||||
class ControlledUnetModel(UNetModel):
|
class ControlledUnetModel(UNetModel):
|
||||||
|
@ -1,6 +1,6 @@
|
|||||||
import torch
|
import torch
|
||||||
from torch import nn, einsum
|
from torch import nn, einsum
|
||||||
from ldm.modules.attention import CrossAttention
|
from .ldm.modules.attention import CrossAttention
|
||||||
from inspect import isfunction
|
from inspect import isfunction
|
||||||
|
|
||||||
|
|
||||||
|
@ -3,11 +3,11 @@ import torch
|
|||||||
import torch.nn.functional as F
|
import torch.nn.functional as F
|
||||||
from contextlib import contextmanager
|
from contextlib import contextmanager
|
||||||
|
|
||||||
from ldm.modules.diffusionmodules.model import Encoder, Decoder
|
from comfy.ldm.modules.diffusionmodules.model import Encoder, Decoder
|
||||||
from ldm.modules.distributions.distributions import DiagonalGaussianDistribution
|
from comfy.ldm.modules.distributions.distributions import DiagonalGaussianDistribution
|
||||||
|
|
||||||
from ldm.util import instantiate_from_config
|
from comfy.ldm.util import instantiate_from_config
|
||||||
from ldm.modules.ema import LitEma
|
from comfy.ldm.modules.ema import LitEma
|
||||||
|
|
||||||
# class AutoencoderKL(pl.LightningModule):
|
# class AutoencoderKL(pl.LightningModule):
|
||||||
class AutoencoderKL(torch.nn.Module):
|
class AutoencoderKL(torch.nn.Module):
|
||||||
|
@ -4,7 +4,7 @@ import torch
|
|||||||
import numpy as np
|
import numpy as np
|
||||||
from tqdm import tqdm
|
from tqdm import tqdm
|
||||||
|
|
||||||
from ldm.modules.diffusionmodules.util import make_ddim_sampling_parameters, make_ddim_timesteps, noise_like, extract_into_tensor
|
from comfy.ldm.modules.diffusionmodules.util import make_ddim_sampling_parameters, make_ddim_timesteps, noise_like, extract_into_tensor
|
||||||
|
|
||||||
|
|
||||||
class DDIMSampler(object):
|
class DDIMSampler(object):
|
||||||
|
@ -19,12 +19,12 @@ from tqdm import tqdm
|
|||||||
from torchvision.utils import make_grid
|
from torchvision.utils import make_grid
|
||||||
# from pytorch_lightning.utilities.distributed import rank_zero_only
|
# from pytorch_lightning.utilities.distributed import rank_zero_only
|
||||||
|
|
||||||
from ldm.util import log_txt_as_img, exists, default, ismap, isimage, mean_flat, count_params, instantiate_from_config
|
from comfy.ldm.util import log_txt_as_img, exists, default, ismap, isimage, mean_flat, count_params, instantiate_from_config
|
||||||
from ldm.modules.ema import LitEma
|
from comfy.ldm.modules.ema import LitEma
|
||||||
from ldm.modules.distributions.distributions import normal_kl, DiagonalGaussianDistribution
|
from comfy.ldm.modules.distributions.distributions import normal_kl, DiagonalGaussianDistribution
|
||||||
from ldm.models.autoencoder import IdentityFirstStage, AutoencoderKL
|
from ..autoencoder import IdentityFirstStage, AutoencoderKL
|
||||||
from ldm.modules.diffusionmodules.util import make_beta_schedule, extract_into_tensor, noise_like
|
from comfy.ldm.modules.diffusionmodules.util import make_beta_schedule, extract_into_tensor, noise_like
|
||||||
from ldm.models.diffusion.ddim import DDIMSampler
|
from .ddim import DDIMSampler
|
||||||
|
|
||||||
|
|
||||||
__conditioning_keys__ = {'concat': 'c_concat',
|
__conditioning_keys__ = {'concat': 'c_concat',
|
||||||
|
@ -6,7 +6,7 @@ from torch import nn, einsum
|
|||||||
from einops import rearrange, repeat
|
from einops import rearrange, repeat
|
||||||
from typing import Optional, Any
|
from typing import Optional, Any
|
||||||
|
|
||||||
from ldm.modules.diffusionmodules.util import checkpoint
|
from .diffusionmodules.util import checkpoint
|
||||||
from .sub_quadratic_attention import efficient_dot_product_attention
|
from .sub_quadratic_attention import efficient_dot_product_attention
|
||||||
|
|
||||||
from comfy import model_management
|
from comfy import model_management
|
||||||
@ -21,7 +21,7 @@ if model_management.xformers_enabled():
|
|||||||
import os
|
import os
|
||||||
_ATTN_PRECISION = os.environ.get("ATTN_PRECISION", "fp32")
|
_ATTN_PRECISION = os.environ.get("ATTN_PRECISION", "fp32")
|
||||||
|
|
||||||
from cli_args import args
|
from comfy.cli_args import args
|
||||||
|
|
||||||
def exists(val):
|
def exists(val):
|
||||||
return val is not None
|
return val is not None
|
||||||
|
@ -6,7 +6,7 @@ import numpy as np
|
|||||||
from einops import rearrange
|
from einops import rearrange
|
||||||
from typing import Optional, Any
|
from typing import Optional, Any
|
||||||
|
|
||||||
from ldm.modules.attention import MemoryEfficientCrossAttention
|
from ..attention import MemoryEfficientCrossAttention
|
||||||
from comfy import model_management
|
from comfy import model_management
|
||||||
|
|
||||||
if model_management.xformers_enabled_vae():
|
if model_management.xformers_enabled_vae():
|
||||||
|
@ -6,7 +6,7 @@ import torch as th
|
|||||||
import torch.nn as nn
|
import torch.nn as nn
|
||||||
import torch.nn.functional as F
|
import torch.nn.functional as F
|
||||||
|
|
||||||
from ldm.modules.diffusionmodules.util import (
|
from .util import (
|
||||||
checkpoint,
|
checkpoint,
|
||||||
conv_nd,
|
conv_nd,
|
||||||
linear,
|
linear,
|
||||||
@ -15,8 +15,8 @@ from ldm.modules.diffusionmodules.util import (
|
|||||||
normalization,
|
normalization,
|
||||||
timestep_embedding,
|
timestep_embedding,
|
||||||
)
|
)
|
||||||
from ldm.modules.attention import SpatialTransformer
|
from ..attention import SpatialTransformer
|
||||||
from ldm.util import exists
|
from comfy.ldm.util import exists
|
||||||
|
|
||||||
|
|
||||||
# dummy replace
|
# dummy replace
|
||||||
|
@ -3,8 +3,8 @@ import torch.nn as nn
|
|||||||
import numpy as np
|
import numpy as np
|
||||||
from functools import partial
|
from functools import partial
|
||||||
|
|
||||||
from ldm.modules.diffusionmodules.util import extract_into_tensor, make_beta_schedule
|
from .util import extract_into_tensor, make_beta_schedule
|
||||||
from ldm.util import default
|
from comfy.ldm.util import default
|
||||||
|
|
||||||
|
|
||||||
class AbstractLowScaleModel(nn.Module):
|
class AbstractLowScaleModel(nn.Module):
|
||||||
|
@ -15,7 +15,7 @@ import torch.nn as nn
|
|||||||
import numpy as np
|
import numpy as np
|
||||||
from einops import repeat
|
from einops import repeat
|
||||||
|
|
||||||
from ldm.util import instantiate_from_config
|
from comfy.ldm.util import instantiate_from_config
|
||||||
|
|
||||||
|
|
||||||
def make_beta_schedule(schedule, n_timestep, linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3):
|
def make_beta_schedule(schedule, n_timestep, linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3):
|
||||||
|
@ -1,5 +1,5 @@
|
|||||||
from ldm.modules.diffusionmodules.upscaling import ImageConcatWithNoiseAugmentation
|
from ..diffusionmodules.upscaling import ImageConcatWithNoiseAugmentation
|
||||||
from ldm.modules.diffusionmodules.openaimodel import Timestep
|
from ..diffusionmodules.openaimodel import Timestep
|
||||||
import torch
|
import torch
|
||||||
|
|
||||||
class CLIPEmbeddingNoiseAugmentation(ImageConcatWithNoiseAugmentation):
|
class CLIPEmbeddingNoiseAugmentation(ImageConcatWithNoiseAugmentation):
|
||||||
|
@ -1,6 +1,6 @@
|
|||||||
import psutil
|
import psutil
|
||||||
from enum import Enum
|
from enum import Enum
|
||||||
from cli_args import args
|
from .cli_args import args
|
||||||
|
|
||||||
class VRAMState(Enum):
|
class VRAMState(Enum):
|
||||||
CPU = 0
|
CPU = 0
|
||||||
|
26
comfy/sd.py
26
comfy/sd.py
@ -2,8 +2,8 @@ import torch
|
|||||||
import contextlib
|
import contextlib
|
||||||
import copy
|
import copy
|
||||||
|
|
||||||
import sd1_clip
|
from . import sd1_clip
|
||||||
import sd2_clip
|
from . import sd2_clip
|
||||||
from comfy import model_management
|
from comfy import model_management
|
||||||
from .ldm.util import instantiate_from_config
|
from .ldm.util import instantiate_from_config
|
||||||
from .ldm.models.autoencoder import AutoencoderKL
|
from .ldm.models.autoencoder import AutoencoderKL
|
||||||
@ -446,10 +446,10 @@ class CLIP:
|
|||||||
else:
|
else:
|
||||||
params = {}
|
params = {}
|
||||||
|
|
||||||
if self.target_clip == "ldm.modules.encoders.modules.FrozenOpenCLIPEmbedder":
|
if self.target_clip.endswith("FrozenOpenCLIPEmbedder"):
|
||||||
clip = sd2_clip.SD2ClipModel
|
clip = sd2_clip.SD2ClipModel
|
||||||
tokenizer = sd2_clip.SD2Tokenizer
|
tokenizer = sd2_clip.SD2Tokenizer
|
||||||
elif self.target_clip == "ldm.modules.encoders.modules.FrozenCLIPEmbedder":
|
elif self.target_clip.endswith("FrozenCLIPEmbedder"):
|
||||||
clip = sd1_clip.SD1ClipModel
|
clip = sd1_clip.SD1ClipModel
|
||||||
tokenizer = sd1_clip.SD1Tokenizer
|
tokenizer = sd1_clip.SD1Tokenizer
|
||||||
|
|
||||||
@ -896,9 +896,9 @@ def load_clip(ckpt_path, embedding_directory=None):
|
|||||||
clip_data = utils.load_torch_file(ckpt_path)
|
clip_data = utils.load_torch_file(ckpt_path)
|
||||||
config = {}
|
config = {}
|
||||||
if "text_model.encoder.layers.22.mlp.fc1.weight" in clip_data:
|
if "text_model.encoder.layers.22.mlp.fc1.weight" in clip_data:
|
||||||
config['target'] = 'ldm.modules.encoders.modules.FrozenOpenCLIPEmbedder'
|
config['target'] = 'comfy.ldm.modules.encoders.modules.FrozenOpenCLIPEmbedder'
|
||||||
else:
|
else:
|
||||||
config['target'] = 'ldm.modules.encoders.modules.FrozenCLIPEmbedder'
|
config['target'] = 'comfy.ldm.modules.encoders.modules.FrozenCLIPEmbedder'
|
||||||
clip = CLIP(config=config, embedding_directory=embedding_directory)
|
clip = CLIP(config=config, embedding_directory=embedding_directory)
|
||||||
clip.load_from_state_dict(clip_data)
|
clip.load_from_state_dict(clip_data)
|
||||||
return clip
|
return clip
|
||||||
@ -974,9 +974,9 @@ def load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, o
|
|||||||
if output_clip:
|
if output_clip:
|
||||||
clip_config = {}
|
clip_config = {}
|
||||||
if "cond_stage_model.model.transformer.resblocks.22.attn.out_proj.weight" in sd_keys:
|
if "cond_stage_model.model.transformer.resblocks.22.attn.out_proj.weight" in sd_keys:
|
||||||
clip_config['target'] = 'ldm.modules.encoders.modules.FrozenOpenCLIPEmbedder'
|
clip_config['target'] = 'comfy.ldm.modules.encoders.modules.FrozenOpenCLIPEmbedder'
|
||||||
else:
|
else:
|
||||||
clip_config['target'] = 'ldm.modules.encoders.modules.FrozenCLIPEmbedder'
|
clip_config['target'] = 'comfy.ldm.modules.encoders.modules.FrozenCLIPEmbedder'
|
||||||
clip = CLIP(config=clip_config, embedding_directory=embedding_directory)
|
clip = CLIP(config=clip_config, embedding_directory=embedding_directory)
|
||||||
w.cond_stage_model = clip.cond_stage_model
|
w.cond_stage_model = clip.cond_stage_model
|
||||||
load_state_dict_to = [w]
|
load_state_dict_to = [w]
|
||||||
@ -997,7 +997,7 @@ def load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, o
|
|||||||
noise_schedule_config["timesteps"] = sd[noise_aug_key].shape[0]
|
noise_schedule_config["timesteps"] = sd[noise_aug_key].shape[0]
|
||||||
noise_schedule_config["beta_schedule"] = "squaredcos_cap_v2"
|
noise_schedule_config["beta_schedule"] = "squaredcos_cap_v2"
|
||||||
params["noise_schedule_config"] = noise_schedule_config
|
params["noise_schedule_config"] = noise_schedule_config
|
||||||
noise_aug_config['target'] = "ldm.modules.encoders.noise_aug_modules.CLIPEmbeddingNoiseAugmentation"
|
noise_aug_config['target'] = "comfy.ldm.modules.encoders.noise_aug_modules.CLIPEmbeddingNoiseAugmentation"
|
||||||
if size == 1280: #h
|
if size == 1280: #h
|
||||||
params["timestep_dim"] = 1024
|
params["timestep_dim"] = 1024
|
||||||
elif size == 1024: #l
|
elif size == 1024: #l
|
||||||
@ -1049,19 +1049,19 @@ def load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, o
|
|||||||
unet_config["in_channels"] = sd['model.diffusion_model.input_blocks.0.0.weight'].shape[1]
|
unet_config["in_channels"] = sd['model.diffusion_model.input_blocks.0.0.weight'].shape[1]
|
||||||
unet_config["context_dim"] = sd['model.diffusion_model.input_blocks.1.1.transformer_blocks.0.attn2.to_k.weight'].shape[1]
|
unet_config["context_dim"] = sd['model.diffusion_model.input_blocks.1.1.transformer_blocks.0.attn2.to_k.weight'].shape[1]
|
||||||
|
|
||||||
sd_config["unet_config"] = {"target": "ldm.modules.diffusionmodules.openaimodel.UNetModel", "params": unet_config}
|
sd_config["unet_config"] = {"target": "comfy.ldm.modules.diffusionmodules.openaimodel.UNetModel", "params": unet_config}
|
||||||
model_config = {"target": "ldm.models.diffusion.ddpm.LatentDiffusion", "params": sd_config}
|
model_config = {"target": "comfy.ldm.models.diffusion.ddpm.LatentDiffusion", "params": sd_config}
|
||||||
|
|
||||||
if noise_aug_config is not None: #SD2.x unclip model
|
if noise_aug_config is not None: #SD2.x unclip model
|
||||||
sd_config["noise_aug_config"] = noise_aug_config
|
sd_config["noise_aug_config"] = noise_aug_config
|
||||||
sd_config["image_size"] = 96
|
sd_config["image_size"] = 96
|
||||||
sd_config["embedding_dropout"] = 0.25
|
sd_config["embedding_dropout"] = 0.25
|
||||||
sd_config["conditioning_key"] = 'crossattn-adm'
|
sd_config["conditioning_key"] = 'crossattn-adm'
|
||||||
model_config["target"] = "ldm.models.diffusion.ddpm.ImageEmbeddingConditionedLatentDiffusion"
|
model_config["target"] = "comfy.ldm.models.diffusion.ddpm.ImageEmbeddingConditionedLatentDiffusion"
|
||||||
elif unet_config["in_channels"] > 4: #inpainting model
|
elif unet_config["in_channels"] > 4: #inpainting model
|
||||||
sd_config["conditioning_key"] = "hybrid"
|
sd_config["conditioning_key"] = "hybrid"
|
||||||
sd_config["finetune_keys"] = None
|
sd_config["finetune_keys"] = None
|
||||||
model_config["target"] = "ldm.models.diffusion.ddpm.LatentInpaintDiffusion"
|
model_config["target"] = "comfy.ldm.models.diffusion.ddpm.LatentInpaintDiffusion"
|
||||||
else:
|
else:
|
||||||
sd_config["conditioning_key"] = "crossattn"
|
sd_config["conditioning_key"] = "crossattn"
|
||||||
|
|
||||||
|
Loading…
Reference in New Issue
Block a user