Merge branch 'main' into development

- this syncs documentation and code
This commit is contained in:
Lincoln Stein
2022-10-09 14:47:27 -04:00
100 changed files with 1976 additions and 1107 deletions

View File

@ -1,4 +0,0 @@
'''
Initialization file for the ldm.dream.generator package
'''
from .base import Generator

View File

@ -1,4 +0,0 @@
'''
Initialization file for the ldm.dream.restoration package
'''
from .base import Restoration

View File

@ -19,7 +19,7 @@ import cv2
import skimage
from omegaconf import OmegaConf
from ldm.dream.generator.base import downsampling
from ldm.invoke.generator.base import downsampling
from PIL import Image, ImageOps
from torch import nn
from pytorch_lightning import seed_everything, logging
@ -28,33 +28,11 @@ from ldm.util import instantiate_from_config
from ldm.models.diffusion.ddim import DDIMSampler
from ldm.models.diffusion.plms import PLMSSampler
from ldm.models.diffusion.ksampler import KSampler
from ldm.dream.pngwriter import PngWriter
from ldm.dream.args import metadata_from_png
from ldm.dream.image_util import InitImageResizer
from ldm.dream.devices import choose_torch_device, choose_precision
from ldm.dream.conditioning import get_uc_and_c
def fix_func(orig):
if hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
def new_func(*args, **kw):
device = kw.get("device", "mps")
kw["device"]="cpu"
return orig(*args, **kw).to(device)
return new_func
return orig
torch.rand = fix_func(torch.rand)
torch.rand_like = fix_func(torch.rand_like)
torch.randn = fix_func(torch.randn)
torch.randn_like = fix_func(torch.randn_like)
torch.randint = fix_func(torch.randint)
torch.randint_like = fix_func(torch.randint_like)
torch.bernoulli = fix_func(torch.bernoulli)
torch.multinomial = fix_func(torch.multinomial)
from ldm.invoke.pngwriter import PngWriter
from ldm.invoke.args import metadata_from_png
from ldm.invoke.image_util import InitImageResizer
from ldm.invoke.devices import choose_torch_device, choose_precision
from ldm.invoke.conditioning import get_uc_and_c
"""Simplified text to image API for stable diffusion/latent diffusion
@ -142,7 +120,8 @@ class Generate:
config = None,
gfpgan=None,
codeformer=None,
esrgan=None
esrgan=None,
free_gpu_mem=False,
):
models = OmegaConf.load(conf)
mconfig = models[model]
@ -169,6 +148,7 @@ class Generate:
self.gfpgan = gfpgan
self.codeformer = codeformer
self.esrgan = esrgan
self.free_gpu_mem = free_gpu_mem
# Note that in previous versions, there was an option to pass the
# device to Generate(). However the device was then ignored, so
@ -295,9 +275,9 @@ class Generate:
def process_image(image,seed):
image.save(f{'images/seed.png'})
The callback used by the prompt2png() can be found in ldm/dream_util.py. It contains code
to create the requested output directory, select a unique informative name for each image, and
write the prompt into the PNG metadata.
The code used to save images to a directory can be found in ldm/invoke/pngwriter.py.
It contains code to create the requested output directory, select a unique informative
name for each image, and write the prompt into the PNG metadata.
"""
# TODO: convert this into a getattr() loop
steps = steps or self.steps
@ -385,7 +365,8 @@ class Generate:
generator = self._make_txt2img()
generator.set_variation(
self.seed, variation_amount, with_variations)
self.seed, variation_amount, with_variations
)
results = generator.generate(
prompt,
iterations=iterations,
@ -521,7 +502,7 @@ class Generate:
)
elif tool == 'outcrop':
from ldm.dream.restoration.outcrop import Outcrop
from ldm.invoke.restoration.outcrop import Outcrop
extend_instructions = {}
for direction,pixels in _pairwise(opt.outcrop):
extend_instructions[direction]=int(pixels)
@ -558,7 +539,7 @@ class Generate:
image_callback = callback,
)
elif tool == 'outpaint':
from ldm.dream.restoration.outpaint import Outpaint
from ldm.invoke.restoration.outpaint import Outpaint
restorer = Outpaint(image,self)
return restorer.process(
opt,
@ -594,18 +575,14 @@ class Generate:
height,
)
if image.width < self.width and image.height < self.height:
print(f'>> WARNING: img2img and inpainting may produce unexpected results with initial images smaller than {self.width}x{self.height} in both dimensions')
# if image has a transparent area and no mask was provided, then try to generate mask
if self._has_transparency(image) and not mask:
print(
'>> Initial image has transparent areas. Will inpaint in these regions.')
if self._check_for_erasure(image):
print(
'>> WARNING: Colors underneath the transparent region seem to have been erased.\n',
'>> Inpainting will be suboptimal. Please preserve the colors when making\n',
'>> a transparency mask, or provide mask explicitly using --init_mask (-M).'
)
if self._has_transparency(image):
self._transparency_check_and_warning(image, mask)
# this returns a torch tensor
init_mask = self._create_init_mask(image,width,height,fit=fit)
init_mask = self._create_init_mask(image, width, height, fit=fit)
if (image.width * image.height) > (self.width * self.height):
print(">> This input is larger than your defaults. If you run out of memory, please use a smaller image.")
@ -621,39 +598,39 @@ class Generate:
def _make_base(self):
if not self.generators.get('base'):
from ldm.dream.generator import Generator
from ldm.invoke.generator import Generator
self.generators['base'] = Generator(self.model, self.precision)
return self.generators['base']
def _make_img2img(self):
if not self.generators.get('img2img'):
from ldm.dream.generator.img2img import Img2Img
from ldm.invoke.generator.img2img import Img2Img
self.generators['img2img'] = Img2Img(self.model, self.precision)
return self.generators['img2img']
def _make_embiggen(self):
if not self.generators.get('embiggen'):
from ldm.dream.generator.embiggen import Embiggen
from ldm.invoke.generator.embiggen import Embiggen
self.generators['embiggen'] = Embiggen(self.model, self.precision)
return self.generators['embiggen']
def _make_txt2img(self):
if not self.generators.get('txt2img'):
from ldm.dream.generator.txt2img import Txt2Img
from ldm.invoke.generator.txt2img import Txt2Img
self.generators['txt2img'] = Txt2Img(self.model, self.precision)
self.generators['txt2img'].free_gpu_mem = self.free_gpu_mem
return self.generators['txt2img']
def _make_txt2img2img(self):
if not self.generators.get('txt2img2'):
from ldm.dream.generator.txt2img2img import Txt2Img2Img
from ldm.invoke.generator.txt2img2img import Txt2Img2Img
self.generators['txt2img2'] = Txt2Img2Img(self.model, self.precision)
self.generators['txt2img2'].free_gpu_mem = self.free_gpu_mem
return self.generators['txt2img2']
def _make_inpaint(self):
if not self.generators.get('inpaint'):
from ldm.dream.generator.inpaint import Inpaint
from ldm.invoke.generator.inpaint import Inpaint
self.generators['inpaint'] = Inpaint(self.model, self.precision)
return self.generators['inpaint']
@ -784,7 +761,7 @@ class Generate:
print(msg)
# Be warned: config is the path to the model config file, not the dream conf file!
# Be warned: config is the path to the model config file, not the invoke conf file!
# Also note that we can get config and weights from self, so why do we need to
# pass them as args?
def _load_model_from_config(self, config, weights):
@ -920,6 +897,17 @@ class Generate:
colored += 1
return colored == 0
def _transparency_check_and_warning(self,image, mask):
if not mask:
print(
'>> Initial image has transparent areas. Will inpaint in these regions.')
if self._check_for_erasure(image):
print(
'>> WARNING: Colors underneath the transparent region seem to have been erased.\n',
'>> Inpainting will be suboptimal. Please preserve the colors when making\n',
'>> a transparency mask, or provide mask explicitly using --init_mask (-M).'
)
def _squeeze_image(self, image):
x, y, resize_needed = self._resolution_check(image.width, image.height)
if resize_needed:

View File

@ -1,7 +1,7 @@
"""Helper class for dealing with image generation arguments.
The Args class parses both the command line (shell) arguments, as well as the
command string passed at the dream> prompt. It serves as the definitive repository
command string passed at the invoke> prompt. It serves as the definitive repository
of all the arguments used by Generate and their default values, and implements the
preliminary metadata standards discussed here:
@ -19,7 +19,7 @@ To use:
print('oops')
sys.exit(-1)
# read in a command passed to the dream> prompt:
# read in a command passed to the invoke> prompt:
opts = opt.parse_cmd('do androids dream of electric sheep? -H256 -W1024 -n4')
# The Args object acts like a namespace object
@ -64,7 +64,7 @@ To generate a dict representing RFC266 metadata:
This will generate an RFC266 dictionary that can then be turned into a JSON
and written to the PNG file. The optional seeds, weights, model_hash and
postprocesser arguments are not available to the opt object and so must be
provided externally. See how dream.py does it.
provided externally. See how invoke.py does it.
Note that this function was originally called format_metadata() and a wrapper
is provided that issues a deprecation notice.
@ -91,8 +91,8 @@ import re
import copy
import base64
import functools
import ldm.dream.pngwriter
from ldm.dream.conditioning import split_weighted_subprompts
import ldm.invoke.pngwriter
from ldm.invoke.conditioning import split_weighted_subprompts
SAMPLER_CHOICES = [
'ddim',
@ -151,7 +151,7 @@ class Args(object):
'''
Initialize new Args class. It takes two optional arguments, an argparse
parser for switches given on the shell command line, and an argparse
parser for switches given on the dream> CLI line. If one or both are
parser for switches given on the invoke> CLI line. If one or both are
missing, it creates appropriate parsers internally.
'''
self._arg_parser = arg_parser or self._create_arg_parser()
@ -168,7 +168,7 @@ class Args(object):
return None
def parse_cmd(self,cmd_string):
'''Parse a dream>-style command string '''
'''Parse a invoke>-style command string '''
command = cmd_string.replace("'", "\\'")
try:
elements = shlex.split(command)
@ -269,7 +269,7 @@ class Args(object):
if a['with_variations']:
formatted_variations = ','.join(f'{seed}:{weight}' for seed, weight in (a["with_variations"]))
switches.append(f'-V {formatted_variations}')
if 'variations' in a:
if 'variations' in a and len(a['variations'])>0:
switches.append(f'-V {a["variations"]}')
return ' '.join(switches)
@ -509,23 +509,23 @@ class Args(object):
)
return parser
# This creates the parser that processes commands on the dream> command line
# This creates the parser that processes commands on the invoke> command line
def _create_dream_cmd_parser(self):
parser = PagingArgumentParser(
formatter_class=ArgFormatter,
description=
"""
*Image generation:*
dream> a fantastic alien landscape -W576 -H512 -s60 -n4
invoke> a fantastic alien landscape -W576 -H512 -s60 -n4
*postprocessing*
!fix applies upscaling/facefixing to a previously-generated image.
dream> !fix 0000045.4829112.png -G1 -U4 -ft codeformer
invoke> !fix 0000045.4829112.png -G1 -U4 -ft codeformer
*History manipulation*
!fetch retrieves the command used to generate an earlier image.
dream> !fetch 0000015.8929913.png
dream> a fantastic alien landscape -W 576 -H 512 -s 60 -A plms -C 7.5
invoke> !fetch 0000015.8929913.png
invoke> a fantastic alien landscape -W 576 -H 512 -s 60 -A plms -C 7.5
!history lists all the commands issued during the current session.
@ -842,7 +842,7 @@ def metadata_from_png(png_file_path) -> Args:
an Args object containing the image metadata. Note that this
returns a single Args object, not multiple.
'''
meta = ldm.dream.pngwriter.retrieve_metadata(png_file_path)
meta = ldm.invoke.pngwriter.retrieve_metadata(png_file_path)
if 'sd-metadata' in meta and len(meta['sd-metadata'])>0 :
return metadata_loads(meta)[0]
else:

View File

@ -0,0 +1,4 @@
'''
Initialization file for the ldm.invoke.generator package
'''
from .base import Generator

View File

@ -1,5 +1,5 @@
'''
Base class for ldm.dream.generator.*
Base class for ldm.invoke.generator.*
including img2img, txt2img, and inpaint
'''
import torch
@ -9,7 +9,7 @@ from tqdm import tqdm, trange
from PIL import Image
from einops import rearrange, repeat
from pytorch_lightning import seed_everything
from ldm.dream.devices import choose_autocast
from ldm.invoke.devices import choose_autocast
from ldm.util import rand_perlin_2d
downsampling = 8
@ -21,6 +21,8 @@ class Generator():
self.seed = None
self.latent_channels = model.channels
self.downsampling_factor = downsampling # BUG: should come from model or config
self.perlin = 0.0
self.threshold = 0
self.variation_amount = 0
self.with_variations = []

View File

@ -1,15 +1,15 @@
'''
ldm.dream.generator.embiggen descends from ldm.dream.generator
and generates with ldm.dream.generator.img2img
ldm.invoke.generator.embiggen descends from ldm.invoke.generator
and generates with ldm.invoke.generator.img2img
'''
import torch
import numpy as np
from tqdm import trange
from PIL import Image
from ldm.dream.generator.base import Generator
from ldm.dream.generator.img2img import Img2Img
from ldm.dream.devices import choose_autocast
from ldm.invoke.generator.base import Generator
from ldm.invoke.generator.img2img import Img2Img
from ldm.invoke.devices import choose_autocast
from ldm.models.diffusion.ddim import DDIMSampler
class Embiggen(Generator):
@ -107,7 +107,7 @@ class Embiggen(Generator):
initsuperwidth = round(initsuperwidth*embiggen[0])
initsuperheight = round(initsuperheight*embiggen[0])
if embiggen[1] > 0: # No point in ESRGAN upscaling if strength is set zero
from ldm.dream.restoration.realesrgan import ESRGAN
from ldm.invoke.restoration.realesrgan import ESRGAN
esrgan = ESRGAN()
print(
f'>> ESRGAN upscaling init image prior to cutting with Embiggen with strength {embiggen[1]}')

View File

@ -1,11 +1,11 @@
'''
ldm.dream.generator.img2img descends from ldm.dream.generator
ldm.invoke.generator.img2img descends from ldm.invoke.generator
'''
import torch
import numpy as np
from ldm.dream.devices import choose_autocast
from ldm.dream.generator.base import Generator
from ldm.invoke.devices import choose_autocast
from ldm.invoke.generator.base import Generator
from ldm.models.diffusion.ddim import DDIMSampler
class Img2Img(Generator):

View File

@ -1,12 +1,12 @@
'''
ldm.dream.generator.inpaint descends from ldm.dream.generator
ldm.invoke.generator.inpaint descends from ldm.invoke.generator
'''
import torch
import numpy as np
from einops import rearrange, repeat
from ldm.dream.devices import choose_autocast
from ldm.dream.generator.img2img import Img2Img
from ldm.invoke.devices import choose_autocast
from ldm.invoke.generator.img2img import Img2Img
from ldm.models.diffusion.ddim import DDIMSampler
from ldm.models.diffusion.ksampler import KSampler
@ -27,7 +27,7 @@ class Inpaint(Img2Img):
# klms samplers not supported yet, so ignore previous sampler
if isinstance(sampler,KSampler):
print(
f">> sampler '{sampler.__class__.__name__}' is not yet supported for inpainting, using DDIMSampler instead."
f">> Using recommended DDIM sampler for inpainting."
)
sampler = DDIMSampler(self.model, device=self.model.device)

View File

@ -1,10 +1,10 @@
'''
ldm.dream.generator.txt2img inherits from ldm.dream.generator
ldm.invoke.generator.txt2img inherits from ldm.invoke.generator
'''
import torch
import numpy as np
from ldm.dream.generator.base import Generator
from ldm.invoke.generator.base import Generator
class Txt2Img(Generator):
def __init__(self, model, precision):

View File

@ -1,11 +1,11 @@
'''
ldm.dream.generator.txt2img inherits from ldm.dream.generator
ldm.invoke.generator.txt2img inherits from ldm.invoke.generator
'''
import torch
import numpy as np
import math
from ldm.dream.generator.base import Generator
from ldm.invoke.generator.base import Generator
from ldm.models.diffusion.ddim import DDIMSampler
@ -64,7 +64,7 @@ class Txt2Img2Img(Generator):
)
print(
f"\n>> Interpolating from {init_width}x{init_height} to {width}x{height}"
f"\n>> Interpolating from {init_width}x{init_height} to {width}x{height} using DDIM sampling"
)
# resizing
@ -75,17 +75,19 @@ class Txt2Img2Img(Generator):
)
t_enc = int(strength * steps)
ddim_sampler = DDIMSampler(self.model, device=self.model.device)
ddim_sampler.make_schedule(
ddim_num_steps=steps, ddim_eta=ddim_eta, verbose=False
)
x = self.get_noise(width,height,False)
z_enc = sampler.stochastic_encode(
z_enc = ddim_sampler.stochastic_encode(
samples,
torch.tensor([t_enc]).to(self.model.device),
noise=x
noise=self.get_noise(width,height,False)
)
# decode it
samples = sampler.decode(
samples = ddim_sampler.decode(
z_enc,
c,
t_enc,

View File

@ -1,17 +1,17 @@
"""
Readline helper functions for dream.py (linux and mac only).
Readline helper functions for invoke.py.
You may import the global singleton `completer` to get access to the
completer object itself. This is useful when you want to autocomplete
seeds:
from ldm.dream.readline import completer
from ldm.invoke.readline import completer
completer.add_seed(18247566)
completer.add_seed(9281839)
"""
import os
import re
import atexit
from ldm.dream.args import Args
from ldm.invoke.args import Args
# ---------------readline utilities---------------------
try:
@ -20,7 +20,7 @@ try:
except (ImportError,ModuleNotFoundError):
readline_available = False
IMG_EXTENSIONS = ('.png','.jpg','.jpeg')
IMG_EXTENSIONS = ('.png','.jpg','.jpeg','.PNG','.JPG','.JPEG','.gif','.GIF')
COMMANDS = (
'--steps','-s',
'--seed','-S',
@ -74,7 +74,7 @@ class Completer(object):
def complete(self, text, state):
'''
Completes dream command line.
Completes invoke command line.
BUG: it doesn't correctly complete files that have spaces in the name.
'''
buffer = readline.get_line_buffer()
@ -287,7 +287,7 @@ def get_completer(opt:Args)->Completer:
readline.parse_and_bind('set skip-completed-text on')
readline.parse_and_bind('set show-all-if-ambiguous on')
histfile = os.path.join(os.path.expanduser(opt.outdir), '.dream_history')
histfile = os.path.join(os.path.expanduser(opt.outdir), '.invoke_history')
try:
readline.read_history_file(histfile)
readline.set_history_length(1000)

View File

@ -0,0 +1,4 @@
'''
Initialization file for the ldm.invoke.restoration package
'''
from .base import Restoration

View File

@ -23,16 +23,16 @@ class Restoration():
# Face Restore Models
def load_gfpgan(self, gfpgan_dir, gfpgan_model_path):
from ldm.dream.restoration.gfpgan import GFPGAN
from ldm.invoke.restoration.gfpgan import GFPGAN
return GFPGAN(gfpgan_dir, gfpgan_model_path)
def load_codeformer(self):
from ldm.dream.restoration.codeformer import CodeFormerRestoration
from ldm.invoke.restoration.codeformer import CodeFormerRestoration
return CodeFormerRestoration()
# Upscale Models
def load_esrgan(self, esrgan_bg_tile=400):
from ldm.dream.restoration.realesrgan import ESRGAN
from ldm.invoke.restoration.realesrgan import ESRGAN
esrgan = ESRGAN(esrgan_bg_tile)
print('>> ESRGAN Initialized')
return esrgan;

View File

@ -8,7 +8,7 @@ pretrained_model_url = 'https://github.com/sczhou/CodeFormer/releases/download/v
class CodeFormerRestoration():
def __init__(self,
codeformer_dir='ldm/dream/restoration/codeformer',
codeformer_dir='ldm/invoke/restoration/codeformer',
codeformer_model_path='weights/codeformer.pth') -> None:
self.model_path = os.path.join(codeformer_dir, codeformer_model_path)
self.codeformer_model_exists = os.path.isfile(self.model_path)
@ -27,7 +27,7 @@ class CodeFormerRestoration():
from basicsr.utils.download_util import load_file_from_url
from basicsr.utils import img2tensor, tensor2img
from facexlib.utils.face_restoration_helper import FaceRestoreHelper
from ldm.dream.restoration.codeformer_arch import CodeFormer
from ldm.invoke.restoration.codeformer_arch import CodeFormer
from torchvision.transforms.functional import normalize
from PIL import Image
@ -35,7 +35,7 @@ class CodeFormerRestoration():
cf = cf_class(dim_embd=512, codebook_size=1024, n_head=8, n_layers=9, connect_list=['32', '64', '128', '256']).to(device)
checkpoint_path = load_file_from_url(url=pretrained_model_url, model_dir=os.path.abspath('ldm/dream/restoration/codeformer/weights'), progress=True)
checkpoint_path = load_file_from_url(url=pretrained_model_url, model_dir=os.path.abspath('ldm/invoke/restoration/codeformer/weights'), progress=True)
checkpoint = torch.load(checkpoint_path)['params_ema']
cf.load_state_dict(checkpoint)
cf.eval()

View File

@ -5,7 +5,7 @@ from torch import nn, Tensor
import torch.nn.functional as F
from typing import Optional, List
from ldm.dream.restoration.vqgan_arch import *
from ldm.invoke.restoration.vqgan_arch import *
from basicsr.utils import get_root_logger
from basicsr.utils.registry import ARCH_REGISTRY

View File

@ -13,8 +13,6 @@ class Outpaint(object):
seed = old_opt.seed
prompt = old_opt.prompt
print(f'DEBUG: old seed={seed}, old prompt = {prompt}')
def wrapped_callback(img,seed,**kwargs):
image_callback(img,seed,use_prefix=prefix,**kwargs)

View File

@ -4,9 +4,9 @@ import copy
import base64
import mimetypes
import os
from ldm.dream.args import Args, metadata_dumps
from ldm.invoke.args import Args, metadata_dumps
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from ldm.dream.pngwriter import PngWriter
from ldm.invoke.pngwriter import PngWriter
from threading import Event
def build_opt(post_data, seed, gfpgan_model_exists):

View File

@ -4,7 +4,7 @@ import torch
import numpy as np
from tqdm import tqdm
from functools import partial
from ldm.dream.devices import choose_torch_device
from ldm.invoke.devices import choose_torch_device
from ldm.models.diffusion.sampler import Sampler
from ldm.modules.diffusionmodules.util import noise_like

View File

@ -2,7 +2,7 @@
import k_diffusion as K
import torch
import torch.nn as nn
from ldm.dream.devices import choose_torch_device
from ldm.invoke.devices import choose_torch_device
from ldm.models.diffusion.sampler import Sampler
from ldm.util import rand_perlin_2d
from ldm.modules.diffusionmodules.util import (
@ -57,8 +57,9 @@ class KSampler(Sampler):
schedule,
steps=model.num_timesteps,
)
self.ds = None
self.s_in = None
self.sigmas = None
self.ds = None
self.s_in = None
def forward(self, x, sigma, uncond, cond, cond_scale):
x_in = torch.cat([x] * 2)
@ -190,7 +191,7 @@ class KSampler(Sampler):
'uncond': unconditional_conditioning,
'cond_scale': unconditional_guidance_scale,
}
print(f'>> Sampling with k_{self.schedule}')
print(f'>> Sampling with k_{self.schedule} starting at step {len(self.sigmas)-S-1} of {len(self.sigmas)-1} ({S} new sampling steps)')
return (
K.sampling.__dict__[f'sample_{self.schedule}'](
model_wrap_cfg, x, sigmas, extra_args=extra_args,
@ -199,6 +200,8 @@ class KSampler(Sampler):
None,
)
# this code will support inpainting if and when ksampler API modified or
# a workaround is found.
@torch.no_grad()
def p_sample(
self,

View File

@ -4,7 +4,7 @@ import torch
import numpy as np
from tqdm import tqdm
from functools import partial
from ldm.dream.devices import choose_torch_device
from ldm.invoke.devices import choose_torch_device
from ldm.models.diffusion.sampler import Sampler
from ldm.modules.diffusionmodules.util import noise_like

View File

@ -8,7 +8,7 @@ import torch
import numpy as np
from tqdm import tqdm
from functools import partial
from ldm.dream.devices import choose_torch_device
from ldm.invoke.devices import choose_torch_device
from ldm.modules.diffusionmodules.util import (
make_ddim_sampling_parameters,

View File

@ -5,7 +5,7 @@ import clip
from einops import rearrange, repeat
from transformers import CLIPTokenizer, CLIPTextModel
import kornia
from ldm.dream.devices import choose_torch_device
from ldm.invoke.devices import choose_torch_device
from ldm.modules.x_transformer import (
Encoder,