Update app.py

app.py

@@ -12,6 +12,8 @@ from Bio.SeqRecord import SeqRecord
 import time
 import urllib.parse
 import base64
+from IPython.display import Image
+import io
 
 urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
 
@@ -40,16 +42,15 @@ def get_protein_image(pdbview):
     pdbview.zoom(0.8)
     pdbview.render()
     png = pdbview.pngImage()
-    print(f"Type of png: {type(png)}")
-    print(f"Dir of png: {dir(png)}")
     
-    # Try to convert to bytes if possible
-    if hasattr(png, 'tobytes'):
-        png_bytes = png.tobytes()
-    elif hasattr(png, 'tostring'):
-        png_bytes = png.tostring()
-    else:
-        raise TypeError(f"Unable to convert {type(png)} to bytes")
+    # Convert to IPython Image
+    img = Image(png)
+    
+    # Save to BytesIO
+    img_io = io.BytesIO()
+    img.save(img_io, format='PNG')
+    img_io.seek(0)
+    png_bytes = img_io.getvalue()
     
     return base64.b64encode(png_bytes).decode('utf-8')