Add inference script

inference_drone_vla.py

@@ -0,0 +1,80 @@
+#!/usr/bin/env python3
+"""
+🎯 Drone VLA Inference: Image + Text → [Vx, Vy, Vz, Yaw_rate]
+
+Usage:
+    python inference_drone_vla.py --model_repo YOUR_USER/drone-vla --image frame.jpg --instruction "fly forward"
+    python inference_drone_vla.py --model_repo YOUR_USER/drone-vla --live --fps 10
+"""
+import os, sys, json, time, argparse, torch, torch.nn as nn, numpy as np
+from PIL import Image
+from transformers import AutoProcessor, AutoModelForImageTextToText
+from peft import PeftModel
+
+class ActionHead(nn.Module):
+    def __init__(self, hidden_size=1024, action_dim=4, hidden_dim=256):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(hidden_size, hidden_dim), nn.GELU(), nn.Dropout(0.1),
+            nn.Linear(hidden_dim, hidden_dim), nn.GELU(), nn.Dropout(0.1),
+            nn.Linear(hidden_dim, action_dim), nn.Tanh())
+    def forward(self, x): return self.net(x)
+
+class DroneVLAInference:
+    def __init__(self, model_repo, base_model="Qwen/Qwen3.5-0.8B", device="cuda"):
+        self.device = torch.device(device if torch.cuda.is_available() else "cpu")
+        print(f"Loading {base_model}...")
+        base = AutoModelForImageTextToText.from_pretrained(base_model, torch_dtype=torch.bfloat16, trust_remote_code=True)
+        self.model = PeftModel.from_pretrained(base, model_repo).to(self.device).eval()
+        from huggingface_hub import hf_hub_download
+        with open(hf_hub_download(model_repo, "vla_config.json")) as f: cfg = json.load(f)
+        self.action_head = ActionHead(cfg.get('hidden_size',1024), cfg.get('action_dim',4))
+        self.action_head.load_state_dict(torch.load(hf_hub_download(model_repo,"action_head.pt"), map_location='cpu'))
+        self.action_head = self.action_head.to(self.device).eval()
+        self.processor = AutoProcessor.from_pretrained(base_model, trust_remote_code=True, min_pixels=224*224, max_pixels=512*512)
+        self.labels = cfg.get('action_labels', ['Vx','Vy','Vz','Yaw_rate'])
+        print(f"✅ Ready! Actions: {self.labels}")
+
+    @torch.no_grad()
+    def predict(self, image, instruction="navigate forward"):
+        if isinstance(image, str): image = Image.open(image).convert('RGB')
+        elif isinstance(image, np.ndarray): image = Image.fromarray(image).convert('RGB')
+        msgs = [{"role":"user","content":[{"type":"image","image":image},{"type":"text","text":instruction}]}]
+        text = self.processor.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
+        inputs = self.processor(text=[text], images=[image], return_tensors="pt").to(self.device)
+        out = self.model(**inputs, output_hidden_states=True, return_dict=True)
+        feat = out.hidden_states[-1][:, -1, :]
+        action = self.action_head(feat.float()).squeeze(0).cpu().numpy()
+        return {'action': action.tolist(), 'labels': {l: float(action[i]) for i,l in enumerate(self.labels)}}
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument('--model_repo', type=str, required=True)
+    p.add_argument('--base_model', type=str, default='Qwen/Qwen3.5-0.8B')
+    p.add_argument('--image', type=str, default=None)
+    p.add_argument('--instruction', type=str, default='navigate forward')
+    p.add_argument('--live', action='store_true')
+    p.add_argument('--fps', type=int, default=10)
+    p.add_argument('--device', type=str, default='cuda')
+    a = p.parse_args()
+    vla = DroneVLAInference(a.model_repo, a.base_model, a.device)
+    
+    if a.image:
+        r = vla.predict(a.image, a.instruction)
+        print(f"\n🎯 Action:")
+        for l, v in r['labels'].items():
+            print(f"  {l:>10s}: {'█'*int(abs(v)*20):<20s} {v:+.3f}")
+    elif a.live:
+        import mss
+        print(f"🔴 Live @ {a.fps}Hz | {a.instruction}")
+        with mss.mss() as sct:
+            while True:
+                t0 = time.time()
+                ss = sct.grab(sct.monitors[1])
+                img = Image.frombytes("RGB", ss.size, ss.bgra, "raw", "BGRX").resize((224,224))
+                r = vla.predict(img, a.instruction)
+                print(f"\r🎮 {' | '.join(f'{k}={v:+.2f}' for k,v in r['labels'].items())}", end='', flush=True)
+                dt = 1/a.fps - (time.time()-t0)
+                if dt > 0: time.sleep(dt)
+
+if __name__ == '__main__': main()