Update app.py

app.py

@@ -17,10 +17,22 @@ model = AutoModelForCausalLM.from_pretrained(model_id, device_map="cuda", trust_
 processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True) 
 model.to("cuda:0")
 
-def base64_to_image(base64_str):
-    img_bytes = base64.b64decode(base64_str)
+from PIL import Image
+import base64
+import zlib
+from io import BytesIO
+
+def decode_and_decompress_image(base64_string):
+    # Decode the Base64 string to bytes
+    compressed_data = base64.b64decode(base64_string.encode('utf-8'))
+
+    # Decompress the data using zlib
+    img_bytes = zlib.decompress(compressed_data)
+
+    # Open the image from bytes
     img_buffer = BytesIO(img_bytes)
     image = Image.open(img_buffer)
+    
     return image
 
 PLACEHOLDER = """
@@ -77,7 +89,7 @@ def bot_streaming(message, history):
     print(f"prompt is -\n{conversation}")
     prompt = processor.tokenizer.apply_chat_template(conversation, tokenize=False, add_generation_prompt=True)
     # image = Image.open(image)
-    image = base64_to_image(image)
+    image = decode_and_decompress_image(image)
     inputs = processor(prompt, image, return_tensors="pt").to("cuda:0") 
 
     streamer = TextIteratorStreamer(processor, **{"skip_special_tokens": True, "skip_prompt": True, 'clean_up_tokenization_spaces':False,})