tanner/min-dalle-test
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

control super condition factor

Browse Source
This commit is contained in:
Brett Kuprel
2022-07-04 21:30:27 -04:00
parent f766a33dab
commit 6d81d514a0
5 changed files with 18 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
min_dalle/models/dalle_bart_decoder.py
View File

@@ -116,7 +116,6 @@ class DalleBartDecoder(nn.Module):
super().__init__()
self.layer_count = layer_count
self.embed_count = embed_count
self.condition_factor = 10.0
self.embed_tokens = nn.Embedding(image_vocab_count + 1, embed_count)
self.embed_positions = nn.Embedding(IMAGE_TOKEN_COUNT, embed_count)
self.layers: List[DecoderLayer] = nn.ModuleList([
@@ -141,6 +140,7 @@ class DalleBartDecoder(nn.Module):
def decode_step(
self,
log2_supercondition_factor: int,
attention_mask: BoolTensor,
encoder_state: FloatTensor,
attention_state: FloatTensor,
@@ -164,7 +164,7 @@ class DalleBartDecoder(nn.Module):
)
decoder_state = self.final_ln(decoder_state)
logits = self.lm_head(decoder_state)
a = self.condition_factor
a = log2_supercondition_factor
logits: FloatTensor = (
logits[:image_count, -1] * (1 - a) +
logits[image_count:, -1] * a
@@ -182,6 +182,7 @@ class DalleBartDecoder(nn.Module):
def decode_row(
self,
row_index: int,
log2_supercondition_factor: int,
encoder_state: FloatTensor,
attention_mask: BoolTensor,
attention_state: FloatTensor,
@@ -190,6 +191,7 @@ class DalleBartDecoder(nn.Module):
for col_index in range(16):
i = 16 * row_index + col_index
probs, attention_state = self.decode_step(
log2_supercondition_factor = log2_supercondition_factor,
attention_mask = attention_mask,
encoder_state = encoder_state,
attention_state = attention_state,