Update nano_banana.py

nano_banana.py

@@ -7,18 +7,16 @@ from diffusers.modular_pipelines import (
     OutputParam,
 )
 from PIL import Image
-import google.generativeai as genai
-import os
+from google import genai
+from io import BytesIO
 
+client = genai.Client()
 
 class NanoBanana(ModularPipelineBlocks):
     def __init__(self, model_id="gemini-2.5-flash-image-preview"):
         super().__init__()
-        api_key = os.getenv("GEMINI_API_KEY")
-        if api_key is None:
-            raise ValueError("Must provide an API key for Gemini through the `GEMINI_API_KEY` env variable.")
-        genai.configure(api_key=api_key)
-        self.model = genai.GenerativeModel(model_name=model_id)
+        # Cannot initialize the client because it throws a pickling error.
+        self.model_id = model_id
 
     @property
     def expected_components(self):
@@ -50,28 +48,34 @@ class NanoBanana(ModularPipelineBlocks):
         return [
             OutputParam(
                 "output_image",
-                type_hint=str,
+                type_hint=Image.Image,
                 description="Output image",
             ),
             OutputParam(
                 "old_image",
-                type_hint=str,
+                type_hint=Image.Image,
                 description="Old image (if) provided by the user",
             )
         ]
 
-
     def __call__(self, components, state: PipelineState) -> PipelineState:
         block_state = self.get_block_state(state)
 
         old_image = block_state.image
-        prompt = block_state.state.prompt
+        prompt = block_state.prompt
         contents = [prompt]
         if old_image is not None:
-            contents.expand(old_image)
+            contents.append(old_image)
         
-        output = self.model.generate_content(contents=contents)
-        block_state.output_image = output
+        response = client.models.generate_content(
+            model=self.model_id, contents=contents
+
+        )
+        for part in response.candidates[0].content.parts:
+            if part.text is not None:
+                continue
+            elif part.inline_data is not None:
+                block_state.output_image = Image.open(BytesIO(part.inline_data.data))
         
         if old_image is not None:
             block_state.old_image = old_image