Spaces:
Runtime error
Runtime error
File size: 2,571 Bytes
ce68674 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 |
import torch
def conv_forward(self):
def forward(input_tensor, temb, scale=1.0):
hidden_states = input_tensor
hidden_states = self.norm1(hidden_states)
hidden_states = self.nonlinearity(hidden_states)
# import pdb; pdb.set_trace()
if self.upsample is not None:
# upsample_nearest_nhwc fails with large batch sizes. see https://github.com/huggingface/diffusers/issues/984
if hidden_states.shape[0] >= 64:
input_tensor = input_tensor.contiguous()
hidden_states = hidden_states.contiguous()
input_tensor = self.upsample(input_tensor)
hidden_states = self.upsample(hidden_states)
elif self.downsample is not None:
input_tensor = self.downsample(input_tensor)
hidden_states = self.downsample(hidden_states)
hidden_states = self.conv1(hidden_states)
if temb is not None:
temb = self.time_emb_proj(self.nonlinearity(temb))[:, :, None, None, None].repeat(1, 1, hidden_states.shape[2], 1, 1)
if temb is not None and self.time_embedding_norm == "default":
hidden_states = hidden_states + temb
hidden_states = self.norm2(hidden_states)
if temb is not None and self.time_embedding_norm == "scale_shift":
scale, shift = torch.chunk(temb, 2, dim=1)
hidden_states = hidden_states * (1 + scale) + shift
hidden_states = self.nonlinearity(hidden_states)
hidden_states = self.dropout(hidden_states)
hidden_states = self.conv2(hidden_states)
# record hidden state
self.record_hidden_state = hidden_states
if self.conv_shortcut is not None:
input_tensor = self.conv_shortcut(input_tensor)
output_tensor = (input_tensor + hidden_states) / self.output_scale_factor
return output_tensor
return forward
def get_conv_feat(unet):
hidden_state_dict = dict()
for i in range(len(unet.up_blocks)):
for j in range(len(unet.up_blocks[i].resnets)):
module = unet.up_blocks[i].resnets[j]
module_name = f"up_blocks.{i}.resnets.{j}"
# print(module_name)
hidden_state_dict[module_name] = module.record_hidden_state
return hidden_state_dict
def prep_unet_conv(unet):
for i in range(len(unet.up_blocks)):
for j in range(len(unet.up_blocks[i].resnets)):
module = unet.up_blocks[i].resnets[j]
module.forward = conv_forward(module)
return unet
|