Update app.py

app.py

@@ -24,7 +24,7 @@ import base64
 
 def image2str(image):
     buf = io.BytesIO()
-    image.save(buf, format="WEBP")
+    image.save(buf, format="PNG")
     i_str = base64.b64encode(buf.getvalue()).decode()
     return f'<div style="float:left"><img src="data:image/png;base64, {i_str}"></div>'
 
@@ -46,7 +46,6 @@ VQ_HUB = "BAAI/Emu3-VisionTokenizer"
 
 # Prepare models and processors
 # Emu3-Gen model and processor
-"""
 gen_model = AutoModelForCausalLM.from_pretrained(
     EMU_GEN_HUB,
     device_map="cpu",
@@ -54,7 +53,6 @@ gen_model = AutoModelForCausalLM.from_pretrained(
     attn_implementation="flash_attention_2",
     trust_remote_code=True,
 )
-"""
 
 # Emu3-Chat model and processor
 chat_model = AutoModelForCausalLM.from_pretrained(
@@ -77,11 +75,10 @@ processor = Emu3Processor(
 )
 
 print(device)
-# gen_model.to(device)
+gen_model.to(device)
 chat_model.to(device)
 image_tokenizer.to(device)
 
-"""
 @spaces.GPU(duration=120)
 def generate_image(prompt):
     POSITIVE_PROMPT = " masterpiece, film grained, best quality."
@@ -141,7 +138,6 @@ def generate_image(prompt):
         if isinstance(im, Image.Image):
             return im
     return None
-"""
 
 @spaces.GPU
 def vision_language_understanding(image, text):
@@ -180,19 +176,19 @@ def chat(history, user_input, user_image):
         # Append the user input and response to the history
         history = history + [(image2str(user_image) + "<br>" + user_input, response)]
     else:
-        history = history + [(user_input, "Currently do not support image genration, please provide an valid image.")]
-        """
+        # history = history + [(user_input, "Currently do not support image genration, please provide an valid image.")]
+        # """
         # Use Emu3-Gen for image generation
         generated_image = generate_image(user_input)
         if generated_image is not None:
             # Append the user input and generated image to the history
-            history = history + [(user_input, generated_image)]
+            history = history + [(user_input, image2str(generated_image))]
         else:
             # If image generation failed, respond with an error message
             history = history + [
                 (user_input, "Sorry, I could not generate an image.")
             ]
-        """
+        # """
     return history, history, gr.update(value=None)
 
 def clear_input():