Update README.md

README.md

@@ -36,7 +36,7 @@ import torch
 import matplotlib.pyplot as plt
 import numpy as np
 
-feature_extractor = ImageGPTImageProcessor.from_pretrained('openai/imagegpt-small')
+processor = ImageGPTImageProcessor.from_pretrained('openai/imagegpt-small')
 model = ImageGPTForCausalImageModeling.from_pretrained('openai/imagegpt-small')
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -48,8 +48,8 @@ context = torch.full((batch_size, 1), model.config.vocab_size - 1) #initialize w
 context = torch.tensor(context).to(device)
 output = model.generate(pixel_values=context, max_length=model.config.n_positions + 1, temperature=1.0, do_sample=True, top_k=40)
 
-clusters = feature_extractor.clusters
-n_px = feature_extractor.size
+clusters = processor.clusters
+n_px = processor.size
 
 samples = output[:,1:].cpu().detach().numpy()
 samples_img = [np.reshape(np.rint(127.5 * (clusters[s] + 1.0)), [n_px, n_px, 3]).astype(np.uint8) for s in samples] # convert color cluster tokens back to pixels