Use better clean up

qwen2_inference.py

@@ -73,21 +73,24 @@ def run_inference(input_file, model_path, args):
         videos=video_inputs,
         padding=True,
         return_tensors="pt",
-    ).to("cuda")
+    )
 
     # GPU Memory after input processing
     after_input_dump = (torch.cuda.memory_allocated(), torch.cuda.memory_reserved())
 
-    output_ids = model.generate(**inputs,
-                                max_new_tokens=args['max_length'],
-                                do_sample=True,
-                                top_p=args['top_p'],
-                                top_k=args['top_k'],
-                                use_cache=True,
-                                num_return_sequences=1,
-                                pad_token_id=processor.tokenizer.pad_token_id,
-                                temperature=args['temperature']
-                                )
+    with torch.inference_mode():
+        output_ids = model.generate(
+            **inputs,
+            max_new_tokens=int(args["max_length"]),
+            do_sample=True,
+            top_p=float(args["top_p"]),
+            top_k=int(args["top_k"]),
+            temperature=float(args["temperature"]),
+            use_cache=True,
+            num_return_sequences=1,
+            pad_token_id=processor.tokenizer.pad_token_id,
+        )
+    
     generated_ids = [
         output_ids[len(input_ids):]
         for input_ids, output_ids in zip(inputs.input_ids, output_ids)
@@ -103,4 +106,10 @@ def run_inference(input_file, model_path, args):
     print_gpu_memory("After Input", after_input_dump[0], after_input_dump[1])
     print_gpu_memory("After Generation", after_gen_dump[0], after_gen_dump[1])
 
+    # Clean up
+    del inputs, output_ids, generated_ids, image, image_input, video_inputs
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        torch.cuda.ipc_collect()
+
     return output_text