diff --git a/sqlalchemy_data_model_visualizer.py b/sqlalchemy_data_model_visualizer.py
index 52b2929..1a1b091 100644
--- a/sqlalchemy_data_model_visualizer.py
+++ b/sqlalchemy_data_model_visualizer.py
@@ -1,20 +1,15 @@
 from datetime import datetime
-from typing import Optional
-from enum import Enum
-from decimal import Decimal
-from sqlalchemy.orm import sessionmaker, declarative_base, relationship
-from sqlalchemy import Column, String, DateTime, Integer, Numeric, Boolean, JSON, ForeignKey, LargeBinary, Text, UniqueConstraint, CheckConstraint, text as sql_text
-from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession
+from sqlalchemy.orm import declarative_base, relationship
+from sqlalchemy import Column, String, DateTime, Integer, Numeric, Boolean, JSON, ForeignKey, LargeBinary, Text, CheckConstraint
 from sqlalchemy import inspect
 import graphviz
 from lxml import etree
 import os
-import re
 Base = declarative_base()
 
-def generate_data_model_diagram(models, output_file='my_data_model_diagram', add_labels=True, view_diagram=True):
+def generate_data_model_diagram(models, output_file='my_data_model_diagram', add_labels=True, view_diagram=True, format='svg'):
     # Initialize graph with more advanced visual settings
-    dot = graphviz.Digraph(comment='Interactive Data Models', format='svg', 
+    dot = graphviz.Digraph(comment='Interactive Data Models', format=format, 
                             graph_attr={'bgcolor': '#EEEEEE', 'rankdir': 'TB', 'splines': 'spline'},
                             node_attr={'shape': 'none', 'fontsize': '12', 'fontname': 'Roboto'},
                             edge_attr={'fontsize': '10', 'fontname': 'Roboto'})