3b again

app.py

@@ -54,13 +54,14 @@ with open("bad_words.txt", "r") as f:
 model, image_processor, tokenizer = create_model_and_transforms(
     clip_vision_encoder_pretrained="openai",
     clip_vision_encoder_path="ViT-L-14",
-    lang_encoder_path="anas-awadalla/mpt-7b",
-    tokenizer_path="anas-awadalla/mpt-7b",
-    cross_attn_every_n_layers=4,
+    lang_encoder_path="anas-awadalla/mpt-1b-redpajama-200b-dolly",
+    tokenizer_path="anas-awadalla/mpt-1b-redpajama-200b-dolly",
+    cross_attn_every_n_layers=1,
 )
 
-checkpoint_path = hf_hub_download("openflamingo/OpenFlamingo-9B-vitl-mpt7b", "checkpoint.pt")
+checkpoint_path = hf_hub_download("openflamingo/OpenFlamingo-3B-vitl-mpt1b-langinstruct", "checkpoint.pt")
 model.load_state_dict(torch.load(checkpoint_path), strict=False)
+
 model.eval()
 
 def generate(
@@ -153,13 +154,13 @@ def generate(
     # with torch.cuda.amp.autocast(dtype=torch.bfloat16):
     output = model.generate(
         vision_x=vision_x,
-        lang_x=input_ids.to("cuda"),
-        attention_mask=attention_mask.to("cuda"),
+        lang_x=input_ids,
+        attention_mask=attention_mask,
         max_new_tokens=30,
         num_beams=3,
-        do_sample=True,
-        temperature=0.3,
-        top_k=0,
+        # do_sample=True,
+        # temperature=0.3,
+        # top_k=0,
     )
 
     gen_text = tokenizer.decode(

open_flamingo/open_flamingo/src/factory.py

@@ -1,6 +1,5 @@
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import open_clip
-import torch
 
 from .flamingo import Flamingo
 from .flamingo_lm import FlamingoLMMixin
@@ -58,7 +57,8 @@ def create_model_and_transforms(
     lang_encoder = AutoModelForCausalLM.from_pretrained(
         lang_encoder_path,
         local_files_only=use_local_files,
-        trust_remote_code=True)
+        trust_remote_code=True,
+    )
 
     # hacks for MPT-1B, which doesn't have a get_input_embeddings method
     if "mpt-1b-redpajama-200b" in lang_encoder_path:
@@ -79,7 +79,6 @@ def create_model_and_transforms(
         decoder_layers_attr_name = _infer_decoder_layers_attr_name(lang_encoder)
     lang_encoder.set_decoder_layers_attr_name(decoder_layers_attr_name)
     lang_encoder.resize_token_embeddings(len(text_tokenizer))
-    lang_encoder.to(0)
 
     model = Flamingo(
         vision_encoder,
@@ -90,7 +89,8 @@ def create_model_and_transforms(
             "width"
         ],
         cross_attn_every_n_layers=cross_attn_every_n_layers,
-        **flamingo_kwargs)
+        **flamingo_kwargs,
+    )
 
     # Freeze all parameters
     model.requires_grad_(False)

open_flamingo/open_flamingo/src/flamingo.py

@@ -212,7 +212,7 @@ class Flamingo(nn.Module):
         with torch.no_grad():
             vision_x = self.vision_encoder(vision_x)[1]
         vision_x = rearrange(vision_x, "(b T F) v d -> b T F v d", b=b, T=T, F=F)
-        vision_x = self.perceiver(vision_x).to(0)
+        vision_x = self.perceiver(vision_x)
 
         for layer in self.lang_encoder._get_decoder_layers():
             layer.condition_vis_x(vision_x)