Update app.py

app.py

@@ -1,21 +1,77 @@
+# import os
+# import gradio as gr
+# from google import genai
+# from PIL import Image
+# from io import BytesIO
+
+# # Load API key from Hugging Face secrets
+# api_key = os.getenv("Gemini_api_key")
+
+# # Initialize Gemini client
+# client = genai.Client(api_key=api_key)
+
+# # Function to generate image from prompt
+# def generate_image(prompt):
+#     response = client.models.generate_content(
+#         model="gemini-2.5-flash-image-preview",
+#         contents=[prompt],
+#     )
+#     for part in response.candidates[0].content.parts:
+#         if part.inline_data is not None:
+#             image = Image.open(BytesIO(part.inline_data.data))
+#             return image
+#     return None
+
+# # Gradio interface
+# demo = gr.Interface(
+#     fn=generate_image,
+#     inputs=gr.Textbox(label="Enter your prompt"),
+#     outputs=gr.Image(label="Generated Image"),
+#     title="Image Generator",
+#     description="Enter a text prompt and generate images"
+# )
+
+# if __name__ == "__main__":
+#     demo.launch()
+
+
+
+
 import os
 import gradio as gr
 from google import genai
 from PIL import Image
 from io import BytesIO
 
-# Load API key from Hugging Face secrets
+# Load API key from environment
 api_key = os.getenv("Gemini_api_key")
 
 # Initialize Gemini client
 client = genai.Client(api_key=api_key)
 
-# Function to generate image from prompt
-def generate_image(prompt):
+# Function to generate or edit image
+def generate_or_edit_image(prompt, uploaded_img):
+    contents = [prompt]
+
+    # If user uploads an image, include it in the request
+    if uploaded_img is not None:
+        # Convert uploaded image to bytes (keep original format if possible)
+        buffered = BytesIO()
+        uploaded_img.save(buffered, format="PNG")  # convert to PNG for consistency
+        image_bytes = buffered.getvalue()
+
+        contents.append({
+            "mime_type": "image/png",
+            "data": image_bytes
+        })
+
+    # Call Gemini image model
     response = client.models.generate_content(
         model="gemini-2.5-flash-image-preview",
-        contents=[prompt],
+        contents=contents,
     )
+
+    # Extract image from response
     for part in response.candidates[0].content.parts:
         if part.inline_data is not None:
             image = Image.open(BytesIO(part.inline_data.data))
@@ -24,12 +80,20 @@ def generate_image(prompt):
 
 # Gradio interface
 demo = gr.Interface(
-    fn=generate_image,
-    inputs=gr.Textbox(label="Enter your prompt"),
-    outputs=gr.Image(label="Generated Image"),
-    title="Image Generator",
-    description="Enter a text prompt and generate images"
+    fn=generate_or_edit_image,
+    inputs=[
+        gr.Textbox(label="Enter your prompt"),
+        gr.Image(
+            label="Upload an image (optional)",
+            type="pil",
+            file_types=[".jpg", ".jpeg", ".png"]  # allow JPG, JPEG, PNG
+        )
+    ],
+    outputs=gr.Image(label="Generated/Edited Image"),
+    title="Image Generator & Editor",
+    description="Enter a text prompt to generate a new image, or upload an image (JPG, JPEG, PNG) to edit it."
 )
 
 if __name__ == "__main__":
     demo.launch()
+