ailia-models
ailia-models copied to clipboard
axinc-ai
ADD StableDiffusion
Model is published today. And also check https://github.com/huggingface/diffusers
pipeline https://github.com/mystic-ai/pipeline/blob/paul/stable-diffusion/examples/custom_models/stable_diffusion.py
○ ldm/modules/attention.py 以下のエラー対策
model has the following type (tensor(float16))
Found kernel for Op with name (LayerNormalization) and type (LayerNormalization) in the supported version range (node_version: 1 kernel start version: 1 kernel_end_version: 2147483647). However the types are incompatible. This op has been implemented only for the following types (tensor(double),), but the node in the model has the following type (tensor(float16))
class BasicTransformerBlock(nn.Module):
...
def _forward(self, x, context=None):
x = self.attn1(self.norm1(x)) + x
x = self.attn2(self.norm2(x), context=context) + x
x = self.ff(self.norm3(x)) + x
return x
↓
class BasicTransformerBlock(nn.Module):
...
def _forward(self, x, context=None):
x = x.float()
x = self.attn1(self.norm1(x)) + x
x = x.float()
x = self.attn2(self.norm2(x), context=context) + x
x = x.float()
x = self.ff(self.norm3(x)) + x
return x
@ooe1123 StableDiffusionのエクスポートの手順につきまして、どのリポジトリを使用したかの情報をいただくことは可能でしょうか。また、可能であれば、エクスポートに使用したリポジトリを、axincのリポジトリにpublic forkいただければと考えています。
@kyakuno こちらのリポジトリで、 https://github.com/CompVis/stable-diffusion どのリビジョンだったかはっきりとは分からないのですが、 時期的に考えて以下の時点のコミットなのではないかと思います。
Date: 23 8 2022 01:57:55 +09:00
○ ldm/models/diffusion/ddpm.py
class LatentDiffusion(DDPM):
...
def decode_first_stage(self, z, predict_cids=False, force_not_quantize=False):
...
return self.first_stage_model.decode(z)
...
class DiffusionWrapper(pl.LightningModule):
...
def forward(self, x, t, c_concat: list = None, c_crossattn: list = None, c_adm=None):
...
elif self.conditioning_key == 'crossattn':
...
if hasattr(self, "scripted_diffusion_model"):
...
else:
out = self.diffusion_model(x, t, context=cc)
↓
class LatentDiffusion(DDPM):
...
def decode_first_stage(self, z, predict_cids=False, force_not_quantize=False):
...
print("------>")
self.first_stage_model.forward = self.first_stage_model.decode
from torch.autograd import Variable
x = Variable(z)
torch.onnx.export(
self.first_stage_model, x, 'autoencoder.onnx',
input_names=["input"],
output_names=["output"],
dynamic_axes={'input': {0:'n', 2:'h', 3:'w'}, 'output': {0:'n', 2:'ho', 3:'wo'}},
verbose=False, opset_version=11
)
print("<------")
...
class DiffusionWrapper(pl.LightningModule):
...
def forward(self, x, t, c_concat: list = None, c_crossattn: list = None, c_adm=None):
...
elif self.conditioning_key == 'crossattn':
...
if hasattr(self, "scripted_diffusion_model"):
...
else:
out = self.diffusion_model(x, t, cc)
# print("------>")
# from torch.autograd import Variable
# xx = (Variable(x.type(torch.float16)), Variable(t), Variable(cc))
# torch.onnx.export(
# self.diffusion_model, xx, 'diffusion_emb.onnx',
# input_names=["x", "timesteps", "context"],
# output_names=["h", "emb", "h0", "h1", "h2", "h3", "h4", "h5", "h6", "h7", "h8", "h9", "h10", "h11"],
# dynamic_axes={'x': {0:'n', 2:'h',3:'w'}, 'timesteps': {0:'n'}, 'context': {0:'n', 1:'l'}, 'h': {0:'n', 2:'h1',3:'w1'}, 'emb': {0:'n'}, 'h0': {0:'n', 2:'h1', 3:'w1'}, 'h1': {0:'n', 2:'h1', 3:'w1'}, 'h2': {0:'n', 2:'h1', 3:'w1'}, 'h3': {0:'n', 2:'h2',3:'w2'}, 'h4': {0:'n', 2:'h2', 3:'w2'}, 'h5': {0:'n', 2:'h2', 3:'w2'}, 'h6': {0:'n', 2:'h3', 3:'w3'}, 'h7': {0:'n', 2:'h3', 3:'w3'}, 'h8': {0:'n', 2:'h3', 3:'w3'}, 'h9': {0:'n', 2:'h4', 3:'w4'}, 'h10': {0:'n', 2:'h4', 3:'w4'}, 'h11': {0:'n', 2:'h4', 3:'w4'}},
# verbose=False, opset_version=12
# )
# print("<------")
h = out[0]
emb = out[1]
hs = out[2:]
h = self.diffusion_model.forward2(
h, emb, cc,
hs[6], hs[7], hs[8], hs[9], hs[10], hs[11])
# print("------>")
# from torch.autograd import Variable
# self.diffusion_model.forward = self.diffusion_model.forward2
# xx = (
# Variable(h), Variable(emb), Variable(cc),
# Variable(hs[6]), Variable(hs[7]), Variable(hs[8]), Variable(hs[9]),
# Variable(hs[10]), Variable(hs[11]))
# torch.onnx.export(
# self.diffusion_model, xx, 'diffusion_mid.onnx',
# input_names=[
# "h", "emb", "context", "h6", "h7", "h8", "h9", "h10", "h11"],
# output_names=["out"],
# dynamic_axes={'h': {0:'n', 2:'h4', 3:'w4'}, 'emb': {0:'n'}, 'context': {0:'n', 1:'l'}, 'h6': {0:'n', 2:'h3', 3:'w3'}, 'h7': {0:'n', 2:'h3', 3:'w3'}, 'h8': {0:'n', 2:'h3', 3:'w3'}, 'h9': {0:'n', 2:'h4', 3:'w4'}, 'h10': {0:'n', 2:'h4', 3:'w4'}, 'h11': {0:'n', 2:'h4', 3:'w4'}, 'out': {0:'n', 2:'h2', 3:'w2'}},
# verbose=False, opset_version=12
# )
# print("<------")
out = self.diffusion_model.forward3(
h, emb, cc,
hs[0], hs[1], hs[2], hs[3], hs[4], hs[5])
# print("------>")
# from torch.autograd import Variable
# self.diffusion_model.forward = self.diffusion_model.forward3
# xx = (
# Variable(h), Variable(emb), Variable(cc),
# Variable(hs[0]), Variable(hs[1]), Variable(hs[2]), Variable(hs[3]),
# Variable(hs[4]), Variable(hs[5]))
# torch.onnx.export(
# self.diffusion_model, xx, 'diffusion_out.onnx',
# input_names=[
# "h", "emb", "context", "h0", "h1", "h2", "h3", "h4", "h5"],
# output_names=["out"],
# dynamic_axes={'h': {0:'n', 2:'h2', 3:'w2'}, 'emb': {0:'n'}, 'context': {0:'n', 1:'l'}, 'h0': {0:'n', 2:'h1', 3:'w1'}, 'h1': {0:'n', 2:'h1', 3:'w1'}, 'h2': {0:'n', 2:'h1', 3:'w1'}, 'h3': {0:'n', 2:'h2', 3:'w2'}, 'h4': {0:'n', 2:'h2', 3:'w2'}, 'h5': {0:'n', 2:'h2', 3:'w2'}, 'out': {0:'n', 2:'h', 3:'w'}},
# verbose=False, opset_version=12
# )
# print("<------")
○ ldm/modules/diffusionmodules/openaimodel.py
class UNetModel(nn.Module):
...
def forward(self, x, timesteps=None, context=None, y=None,**kwargs):
...
h = self.middle_block(h, emb, context)
↓
class UNetModel(nn.Module):
...
def forward(self, x, timesteps=None, context=None, y=None,**kwargs):
...
h = self.middle_block(h, emb, context)
return h, emb, hs[0], hs[1], hs[2], hs[3], hs[4], hs[5], hs[6], hs[7], hs[8], hs[9], hs[10], hs[11]
def forward2(self, h, emb, context, h6, h7, h8, h9, h10, h11):
hs = [h6, h7, h8, h9, h10, h11]
for i, module in enumerate(self.output_blocks[:6]):
h = th.cat([h, hs.pop()], dim=1)
h = module(h, emb, context)
return h
def forward3(self, h, emb, context, h0, h1, h2, h3, h4, h5):
hs = [h0, h1, h2, h3, h4, h5]
for i, module in enumerate(self.output_blocks[6:]):
h = th.cat([h, hs.pop()], dim=1)
h = module(h, emb, context)
if self.predict_codebook_ids:
return self.id_predictor(h)
else:
return self.out(h)
○ ldm/modules/diffusionmodules/util.py
def checkpoint(func, inputs, params, flag):
...
if flag:
...
↓
def checkpoint(func, inputs, params, flag):
...
flag = False
if flag:
...