"_" + "".join([chr(ord("A") + n) for n in range(len(colors) - 1)]) + "><"
)
- def __init__(self, height=6, width=8, nb_birds=3, speed=1, nb_iterations=4):
+ def __init__(
+ self,
+ height=6,
+ width=8,
+ nb_birds=3,
+ speed=2,
+ nb_iterations=2,
+ avoid_collision=True,
+ ):
self.height = height
self.width = width
self.nb_birds = nb_birds
self.speed = speed
self.nb_iterations = nb_iterations
+ self.avoid_collision = avoid_collision
def direction_tokens(self):
return self.token_forward, self.token_backward
frame_sequences = []
for _ in tqdm.tqdm(range(nb), dynamic_ncols=True, desc="world generation"):
- result = torch.zeros(
- self.nb_iterations, self.height, self.width, dtype=torch.int64
- )
-
i, j, vi, vj = (
torch.empty(self.nb_birds, dtype=torch.int64),
torch.empty(self.nb_birds, dtype=torch.int64),
torch.empty(self.nb_birds, dtype=torch.int64),
)
+ def collision_okay():
+ if not self.avoid_collision:
+ return True
+
+ count = torch.zeros(self.height, self.width, dtype=torch.int64)
+
+ for n in range(self.nb_birds):
+ count[i[n], j[n]] += 1
+ count[i[n] - vi[n], j[n]] += 1
+ count[i[n], j[n] - vj[n]] += 1
+
+ return count.max() <= 1
+
col = (
torch.randperm(self.colors.size(0) - 1)[: self.nb_birds].sort().values
+ 1
)
- for n in range(self.nb_birds):
+ while True:
while True:
- i[n] = torch.randint(self.height, (1,))
- j[n] = torch.randint(self.width, (1,))
- vm = torch.randint(4, (1,))
- vi[n], vj[n] = (vm % 2) * 2 - 1, (vm // 2) * 2 - 1
- if (
- i[n] - vi[n] >= 0
- and i[n] - vi[n] < self.height
- and j[n] - vj[n] >= 0
- and j[n] - vj[n] < self.width
- ):
+ for n in range(self.nb_birds):
+ while True:
+ i[n] = torch.randint(self.height, (1,))
+ j[n] = torch.randint(self.width, (1,))
+ vm = torch.randint(4, (1,))
+ vi[n], vj[n] = (vm % 2) * 2 - 1, (vm // 2) * 2 - 1
+ if (
+ i[n] - vi[n] >= 0
+ and i[n] - vi[n] < self.height
+ and j[n] - vj[n] >= 0
+ and j[n] - vj[n] < self.width
+ ):
+ break
+
+ if collision_okay():
break
- for l in range(self.nb_iterations):
- for n in range(self.nb_birds):
- c = col[n]
- result[l, i[n], j[n]] = c
- result[l, i[n] - vi[n], j[n]] = c
- result[l, i[n], j[n] - vj[n]] = c
+ result = torch.zeros(
+ self.nb_iterations * self.speed,
+ self.height,
+ self.width,
+ dtype=torch.int64,
+ )
- if (i[n] == 0 and vi[n] == -1) or (
- i[n] == self.height - 1 and vi[n] == 1
- ):
- vi[n] = -vi[n]
+ fine = torch.empty(self.nb_iterations * self.speed)
- if (j[n] == 0 and vj[n] == -1) or (
- j[n] == self.width - 1 and vj[n] == 1
- ):
- vj[n] = -vj[n]
+ t_to_keep = (
+ torch.arange(self.nb_iterations, device=result.device) * self.speed
+ )
- i[n] += vi[n]
- j[n] += vj[n]
+ for l in range(self.nb_iterations * self.speed):
+ fine[l] = collision_okay()
+ for n in range(self.nb_birds):
+ c = col[n]
+ result[l, i[n], j[n]] = c
+ result[l, i[n] - vi[n], j[n]] = c
+ result[l, i[n], j[n] - vj[n]] = c
+
+ if (i[n] == 0 and vi[n] == -1) or (
+ i[n] == self.height - 1 and vi[n] == 1
+ ):
+ vi[n] = -vi[n]
+
+ if (j[n] == 0 and vj[n] == -1) or (
+ j[n] == self.width - 1 and vj[n] == 1
+ ):
+ vj[n] = -vj[n]
+
+ i[n] += vi[n]
+ j[n] += vj[n]
+
+ result = result[t_to_keep]
+ fine = fine[t_to_keep]
+
+ if fine[-1]:
+ break
frame_sequences.append(result)
return frame_sequences
+ ######################################################################
+
+ def generate_prompts_and_answers(self, nb):
+ frame_sequences = self.generate_frame_sequences(nb)
+ prompts = frame_sequences[:, : frame_sequences.size(0) // 2].flatten(1)
+ answers = frame_sequences[:, frame_sequences.size(0) // 2 :].flatten(1)
+ return prompts, answers
+
def generate_token_sequences(self, nb):
frame_sequences = self.generate_frame_sequences(nb)
if __name__ == "__main__":
import time
- sky = Sky(height=6, width=8, speed=1, nb_iterations=4)
+ sky = Sky(height=6, width=8, speed=4, nb_iterations=2)
start_time = time.perf_counter()
- seq = sky.generate_frame_sequences(nb=64)
+ token_sequences = sky.generate_token_sequences(nb=64)
delay = time.perf_counter() - start_time
- print(f"{seq.size(0)/delay:02f} seq/s")
+ print(f"{token_sequences.size(0)/delay:02f} seq/s")
# print(sky.seq2str(seq[:4]))
# m = (torch.rand(seq.size()) < 0.05).long()
# seq = (1 - m) * seq + m * 23
- print(seq.size())
- img = sky.seq2img(seq)
- print(img.size())
+ # print(seq.size())
+ img = sky.seq2img(token_sequences)
+ # print(img.size())
torchvision.utils.save_image(
img.float() / 255.0, "/tmp/world.png", nrow=6, padding=6, pad_value=0
projects / culture.git / blobdiff