added files

app.py

@@ -61,7 +61,7 @@ model, image_processor, tokenizer = create_model_and_transforms(
 
 checkpoint_path = hf_hub_download("openflamingo/OpenFlamingo-9B-vitl-mpt7b", "checkpoint.pt")
 model.load_state_dict(torch.load(checkpoint_path), strict=False)
-model.eval().to("cuda")
+model.eval()
 
 def generate(
     idx,
@@ -152,7 +152,7 @@ def generate(
 
     # with torch.cuda.amp.autocast(dtype=torch.bfloat16):
     output = model.generate(
-        vision_x=vision_x.to("cuda"),
+        vision_x=vision_x,
         lang_x=input_ids.to("cuda"),
         attention_mask=attention_mask.to("cuda"),
         max_new_tokens=30,

open_flamingo/open_flamingo/src/factory.py

@@ -79,6 +79,7 @@ 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,

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)
+        vision_x = self.perceiver(vision_x).to(0)
 
         for layer in self.lang_encoder._get_decoder_layers():
             layer.condition_vis_x(vision_x)