diff --git a/scripts/python-gen.sh b/scripts/python-gen.sh
index 9440899..11519b4 100755
--- a/scripts/python-gen.sh
+++ b/scripts/python-gen.sh
@@ -2,4 +2,5 @@
 
 mkdir -p strafesnet
 cp *.proto strafesnet/
-python3 -m grpc_tools.protoc -I./strafesnet --python_out=./python-grpc/strafesnet --grpc_python_out=./python-grpc/strafesnet  *.proto
\ No newline at end of file
+python3 -m grpc_tools.protoc -I./strafesnet --python_out=./python-grpc/strafesnet --grpc_python_out=./python-grpc/strafesnet  *.proto
+sed -i 's/^\(import.*pb2\)/from . \1/g' strafesnet/*.py
\ No newline at end of file