Update app.py

app.py

@@ -13,6 +13,10 @@ import boto3
 from botocore.exceptions import ClientError
 from PIL import Image
 
+from google import genai
+from google.genai import types
+from io import BytesIO
+
 torch.set_grad_enabled(False)
 
 pipeline = PuLIDPipeline()
@@ -98,6 +102,33 @@ def run(bucket_folder, person_name, req_id, gender='male'):
 
     id_image = image_paths[0]
 
+    ims = []
+
+    #--------------Google Gemini---------------#
+
+    client = genai.Client(api_key='AIzaSyC5UG7dpBbik4ULenbf-xPYvfObnFv0i8E')
+
+    image = Image.open(id_image)
+
+    try:
+        response = client.models.generate_content(
+            model="gemini-2.5-flash-image-preview",
+            contents=[prompt, image],
+        )
+        
+        for part in response.candidates[0].content.parts:
+            if part.text is not None:
+                print(part.text)
+            elif part.inline_data is not None:
+                image = Image.open(BytesIO(part.inline_data.data))
+        ims.append(image)
+    except:
+        print('Error while calling GEMINI')
+
+        
+    #----------------------------------------#
+
+
     if image_paths is not None:
         id_image = resize_numpy_image_long(np.array(Image.open(image_paths[0])), 1024)
         id_embeddings = pipeline.get_id_embedding(id_image)
@@ -111,7 +142,6 @@ def run(bucket_folder, person_name, req_id, gender='male'):
         id_embeddings = None
 
     seed_everything(seed)
-    ims = []
     for _ in range(n_samples):
         img = pipeline.inference(prompt, (1, H, W), neg_prompt, id_embeddings, id_scale, scale, steps)[0]
         ims.append(np.array(img))