Skip to content

Commit

Permalink
seems to be training
Browse files Browse the repository at this point in the history
  • Loading branch information
lucidrains committed Jul 25, 2024
1 parent 5888a54 commit 1f6a491
Show file tree
Hide file tree
Showing 2 changed files with 10 additions and 2 deletions.
10 changes: 9 additions & 1 deletion autoregressive_diffusion_pytorch/autoregressive_diffusion.py
Original file line number Diff line number Diff line change
Expand Up @@ -511,6 +511,12 @@ def forward(

# image wrapper

def normalize_to_neg_one_to_one(img):
return img * 2 - 1

def unnormalize_to_zero_to_one(t):
return (t + 1) * 0.5

class ImageAutoregressiveDiffusion(Module):
def __init__(
self,
Expand Down Expand Up @@ -541,8 +547,10 @@ def __init__(

def sample(self, batch_size = 1):
tokens = self.model.sample(batch_size = batch_size)
return self.to_image(tokens)
images = self.to_image(tokens)
return unnormalize_to_zero_to_one(images)

def forward(self, images):
images = normalize_to_neg_one_to_one(images)
tokens = self.to_tokens(images)
return self.model(tokens)
2 changes: 1 addition & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
[project]
name = "autoregressive-diffusion-pytorch"
version = "0.0.5"
version = "0.0.6"
description = "Autoregressive Diffusion - Pytorch"
authors = [
{ name = "Phil Wang", email = "[email protected]" }
Expand Down

0 comments on commit 1f6a491

Please sign in to comment.