Update app.py

app.py

@@ -8,6 +8,7 @@ import gradio as gr
 import base64
 import mimetypes
 from translatepy import Translator
+from io import BytesIO
 
 from google import genai
 from google.genai import types
@@ -23,14 +24,13 @@ def translate_to_english(text):
         return result.result
     except Exception as e:
         print(f"Translation error: {e}")
-        return text  # Return original text if translation fails
+        return text
 
 def generate(text, file_name, api_key, model="gemini-2.0-flash-exp"):
-    # Initialize client using provided api_key (or fallback to env variable)
     client = genai.Client(api_key=(api_key.strip() if api_key and api_key.strip() != ""
-                                     else os.environ.get("GEMINI_API_KEY")))
+                                 else os.environ.get("GEMINI_API_KEY")))
     
-    files = [ client.files.upload(file=file_name) ]
+    files = [client.files.upload(file=file_name)]
     
     contents = [
         types.Content(
@@ -54,43 +54,23 @@ def generate(text, file_name, api_key, model="gemini-2.0-flash-exp"):
     )
 
     text_response = ""
-    image_path = None
-    # Create a temporary file to potentially store image data.
-    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
-        temp_path = tmp.name
-        for chunk in client.models.generate_content_stream(
-            model=model,
-            contents=contents,
-            config=generate_content_config,
-        ):
-            if not chunk.candidates or not chunk.candidates[0].content or not chunk.candidates[0].content.parts:
-                continue
-            candidate = chunk.candidates[0].content.parts[0]
-            # Check for inline image data
-            if candidate.inline_data:
-                save_binary_file(temp_path, candidate.inline_data.data)
-                print(f"File of mime type {candidate.inline_data.mime_type} saved to: {temp_path} and prompt input: {text}")
-                image_path = temp_path
-                # If an image is found, we assume that is the desired output.
-                break
-            else:
-                # Accumulate text response if no inline_data is present.
-                text_response += chunk.text + "\n"
+    image_data = None
+    for chunk in client.models.generate_content_stream(
+        model=model,
+        contents=contents,
+        config=generate_content_config,
+    ):
+        if not chunk.candidates or not chunk.candidates[0].content or not chunk.candidates[0].content.parts:
+            continue
+        candidate = chunk.candidates[0].content.parts[0]
+        if candidate.inline_data:
+            image_data = candidate.inline_data.data
+            break
+        else:
+            text_response += chunk.text + "\n"
     
     del files
-    return image_path, text_response
-
-def convert_to_png(image_path):
-    """Convert any image format to PNG and return the path"""
-    with Image.open(image_path) as img:
-        # Convert RGBA to RGB if needed
-        if img.mode == 'RGBA':
-            img = img.convert('RGB')
-        # Create a new temp file with PNG extension
-        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
-            output_path = tmp.name
-            img.save(output_path, format="PNG", quality=95)
-            return output_path
+    return image_data, text_response
 
 def process_image_and_prompt(composite_pil, prompt, gemini_api_key):
     try:
@@ -98,32 +78,43 @@ def process_image_and_prompt(composite_pil, prompt, gemini_api_key):
         translated_prompt = translate_to_english(prompt)
         print(f"Original prompt: {prompt}, Translated prompt: {translated_prompt}")
         
-        # Save the composite image to a temporary file as PNG
+        # Save the composite image to memory as PNG
+        img_byte_arr = BytesIO()
+        composite_pil.save(img_byte_arr, format='PNG')
+        img_byte_arr.seek(0)
+        
+        # Create a temporary PNG file
         with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
-            composite_path = tmp.name
-            composite_pil.save(composite_path, format="PNG")
+            tmp.write(img_byte_arr.getvalue())
+            file_name = tmp.name
         
-        file_name = composite_path  
-        input_text = translated_prompt 
-        model = "gemini-2.0-flash-exp" 
-
-        image_path, text_response = generate(text=input_text, file_name=file_name, api_key=gemini_api_key, model=model)
+        image_data, text_response = generate(
+            text=translated_prompt,
+            file_name=file_name,
+            api_key=gemini_api_key,
+            model="gemini-2.0-flash-exp"
+        )
         
-        if image_path:
-            # Convert the output image to PNG format
-            png_path = convert_to_png(image_path)
-            result_img = Image.open(png_path)
-            return [result_img], ""  # Return image in gallery and empty text output.
+        if image_data:
+            # Convert the binary image data to PNG format
+            img = Image.open(BytesIO(image_data))
+            if img.mode == 'RGBA':
+                img = img.convert('RGB')
+            
+            # Save to BytesIO as PNG
+            output_buffer = BytesIO()
+            img.save(output_buffer, format="PNG")
+            output_buffer.seek(0)
+            
+            # Create PIL Image from buffer
+            result_img = Image.open(output_buffer)
+            return [result_img], ""
         else:
-            # Return no image and the text response.
             return None, text_response
     except Exception as e:
-        raise gr.Error(f"Error Getting {e}", duration=5)
-
+        raise gr.Error(f"Error: {str(e)}", duration=5)
 
-# Build a Blocks-based interface with a custom HTML header and CSS
-with gr.Blocks(css_paths="style.css",) as demo:
-    # Custom HTML header with proper class for styling
+with gr.Blocks(css_paths="style.css") as demo:
     gr.HTML(
     """
     <div class="header-container">
@@ -185,14 +176,17 @@ with gr.Blocks(css_paths="style.css",) as demo:
             submit_btn = gr.Button("Generate", elem_classes="generate-btn")
         
         with gr.Column(elem_classes="output-column"):
-            output_gallery = gr.Gallery(label="Generated Outputs (PNG)", elem_classes="output-gallery")
+            output_gallery = gr.Gallery(
+                label="Generated Outputs (PNG)", 
+                elem_classes="output-gallery",
+                format="png"  # Force Gradio to use PNG format
+            )
             output_text = gr.Textbox(
                 label="Gemini Output", 
                 placeholder="Text response will appear here if no image is generated.",
                 elem_classes="output-text"
             )
 
-    # Set up the interaction with two outputs.
     submit_btn.click(
         fn=process_image_and_prompt,
         inputs=[image_input, prompt_input, gemini_api_key],
@@ -214,7 +208,7 @@ with gr.Blocks(css_paths="style.css",) as demo:
     
     gr.Examples(
         examples=examples,
-        inputs=[image_input, prompt_input,],
+        inputs=[image_input, prompt_input],
         elem_id="examples-grid"
     )