Update app.py

app.py

@@ -62,7 +62,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()
+model.eval().to(0, dtype=torch.bfloat16)
 
 def generate(
     idx,
@@ -151,17 +151,17 @@ def generate(
     vision_x = vision_x.unsqueeze(1).unsqueeze(0)
     print(vision_x.shape)
 
-    # with torch.cuda.amp.autocast(dtype=torch.bfloat16):
-    output = model.generate(
-        vision_x=vision_x,
-        lang_x=input_ids,
-        attention_mask=attention_mask,
-        max_new_tokens=30,
-        num_beams=3,
-        # do_sample=True,
-        # temperature=0.3,
-        # top_k=0,
-    )
+    with torch.cuda.amp.autocast(dtype=torch.bfloat16):
+        output = model.generate(
+            vision_x=vision_x.to(0, dtype=torch.bfloat16),
+            lang_x=input_ids.to(0),
+            attention_mask=attention_mask.to(0),
+            max_new_tokens=30,
+            num_beams=3,
+            # do_sample=True,
+            # temperature=0.3,
+            # top_k=0,
+        )
 
     gen_text = tokenizer.decode(
         output[0][len(input_ids[0]):], skip_special_tokens=True