"""
title: Gemeni Genarate Image
author: KwanLeeFrmVi
author_url: https://gist.github.com/kwanLeeFrmVi
gist_url: https://gist.github.com/kwanLeeFrmVi/e2b0bb11c57077d9a46a9da558ace78c
version: 0.0.1
"""
import os
import requests
import json
import base64
from datetime import datetime
from pydantic import BaseModel, Field
from typing import Callable, Any, Optional
class EventEmitter:
def __init__(self, event_emitter: Optional[Callable[[dict], Any]] = None):
self.event_emitter = event_emitter
async def emit(self, description="Unknown State", status="in_progress", done=False):
if self.event_emitter:
await self.event_emitter(
{
"type": "status",
"data": {
"description": description,
"status": status,
"done": done,
},
}
)
class Tools:
class Valves(BaseModel):
# Gemini image settings
GEMINI_API_ENDPOINT: str = Field(
default="https://generativelanguage.googleapis.com/v1beta/models",
description="Base Gemini API URL",
)
GEMINI_API_KEY: str = Field(
default="YOUR_GEMINI_API_KEY", description="Your Google Gemini API key"
)
GEMINI_IMG_MODEL: str = Field(
default="gemini-2.0-flash-exp-image-generation",
description="Gemini model for image gen & editing",
)
def __init__(self):
self.valves = self.Valves()
async def generate_image_markdown(
self,
prompt: str,
__event_emitter__: Optional[Callable[[dict], Any]] = None,
) -> str:
# 1) notify start
await __event_emitter__(
{
"type": "status",
"data": {"description": f"Starting image gen: {prompt}", "done": False},
}
)
# 2) call Gemini
url = (
f"{self.valves.GEMINI_API_ENDPOINT}/"
f"{self.valves.GEMINI_IMG_MODEL}:generateContent"
f"?key={self.valves.GEMINI_API_KEY}"
)
payload = {
"contents": [{"parts": [{"text": prompt}]}],
"generationConfig": {"responseModalities": ["TEXT", "IMAGE"]},
}
try:
resp = requests.post(
url,
headers={"Content-Type": "application/json"},
json=payload,
timeout=60,
)
resp.raise_for_status()
data = resp.json()
# 3) extract image base64
cand = data["candidates"][0]
if "output" in cand and "image" in cand["output"]:
b64 = cand["output"]["image"]["data"]
elif "result" in cand and "image" in cand["result"]:
b64 = cand["result"]["image"]["data"]
elif (
"content" in cand
and isinstance(cand["content"], dict)
and "parts" in cand["content"]
and isinstance(cand["content"]["parts"], list)
and "inlineData" in cand["content"]["parts"][0]
):
b64 = cand["content"]["parts"][0]["inlineData"]["data"]
else:
raise KeyError(f"No image data in candidate keys: {cand.keys()}")
# 4) emit completion status
await __event_emitter__(
{
"type": "status",
"data": {"description": "Image gen complete", "done": True},
}
)
# 5) emit the markdown image
md = f""
await __event_emitter__({"type": "message", "data": {"content": f"{md}"}})
return f"Image successfully generated"
except Exception as e:
err = str(e)
await __event_emitter__(
{
"type": "status",
"data": {"description": f"Gen failed: {err}", "done": True},
}
)
return f"Error during image gen: {err}"
async def edit_image_markdown(
self,
prompt: str,
image_path: str,
__event_emitter__: Optional[Callable[[dict], Any]] = None,
) -> str:
# 1) notify start
await __event_emitter__(
{
"type": "status",
"data": {
"description": f"Starting image edit: {prompt}",
"done": False,
},
}
)
# 2) encode input image
raw = open(image_path, "rb").read()
b64in = base64.b64encode(raw).decode()
# 3) call Gemini
url = (
f"{self.valves.GEMINI_API_ENDPOINT}/"
f"{self.valves.GEMINI_IMG_MODEL}:generateContent"
f"?key={self.valves.GEMINI_API_KEY}"
)
payload = {
"contents": [
{
"parts": [
{"text": prompt},
{"inline_data": {"mime_type": "image/jpeg", "data": b64in}},
]
}
],
"generationConfig": {"responseModalities": ["TEXT", "IMAGE"]},
}
try:
resp = requests.post(
url,
headers={"Content-Type": "application/json"},
json=payload,
timeout=60,
)
resp.raise_for_status()
data = resp.json()
# 4) extract edited image base64
cand = data["candidates"][0]
if "output" in cand and "image" in cand["output"]:
b64out = cand["output"]["image"]["data"]
elif "result" in cand and "image" in cand["result"]:
b64out = cand["result"]["image"]["data"]
elif (
"content" in cand
and isinstance(cand["content"], dict)
and "parts" in cand["content"]
and "inlineData" in cand["content"]["parts"][0]
):
b64out = cand["content"]["parts"][0]["inlineData"]["data"]
else:
raise KeyError(f"No image data in candidate keys: {cand.keys()}")
# 5) emit completion status
await __event_emitter__(
{
"type": "status",
"data": {"description": "Image edit complete", "done": True},
}
)
# 6) emit the markdown image
md = f""
await __event_emitter__({"type": "message", "data": {"content": f"{md}"}})
return f"Image successfully edited"
except Exception as e:
err = str(e)
await __event_emitter__(
{
"type": "status",
"data": {"description": f"Edit failed: {err}", "done": True},
}
)
return f"Error during image edit: {err}"
