Fix imports.

This commit is contained in:
comfyanonymous 2023-05-04 18:07:41 -04:00
parent 7e51bbd07f
commit bae4fb4a9d
13 changed files with 42 additions and 42 deletions

View File

@ -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):

View File

@ -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

View File

@ -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):

View File

@ -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):

View File

@ -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',

View File

@ -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

View File

@ -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():

View File

@ -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

View File

@ -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):

View File

@ -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):

View File

@ -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):

View File

@ -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

View File

@ -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"