Update README.md

README.md

@@ -100,7 +100,8 @@ GENERATION_CONFIG = GenerationConfig(
     top_k=2048,
 )
 
-h, w = pos_inputs.image_size[0]
+h = pos_inputs.image_size[:, 0]
+w = pos_inputs.image_size[:, 1]
 constrained_fn = processor.build_prefix_constrained_fn(h, w)
 logits_processor = LogitsProcessorList([
     UnbatchedClassifierFreeGuidanceLogitsProcessor(
@@ -118,7 +119,8 @@ logits_processor = LogitsProcessorList([
 outputs = model.generate(
     pos_inputs.input_ids.to("cuda:0"),
     GENERATION_CONFIG,
-    logits_processor=logits_processor
+    logits_processor=logits_processor,
+    attention_mask=pos_inputs.attention_mask.to("cuda:0"),
 )
 
 mm_list = processor.decode(outputs[0])
@@ -139,12 +141,16 @@ inputs = processor(
     padding="longest",
     return_tensors="pt",
 )
-GENERATION_CONFIG = GenerationConfig(pad_token_id=tokenizer.pad_token_id, bos_token_id=tokenizer.bos_token_id, eos_token_id=tokenizer.eos_token_id)
+GENERATION_CONFIG = GenerationConfig(
+    pad_token_id=tokenizer.pad_token_id,
+    bos_token_id=tokenizer.bos_token_id,
+    eos_token_id=tokenizer.eos_token_id,
+    max_new_tokens=1024,
+)
 
  outputs = model.generate(
      inputs.input_ids.to("cuda:0"),
      GENERATION_CONFIG,
-     max_new_tokens=1024,
      attention_mask=inputs.attention_mask.to("cuda:0"),
  )
  outputs = outputs[:, inputs.input_ids.shape[-1]:]