Spaces:
Running
Running
File size: 1,891 Bytes
ac3fd77 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 |
import base64
import os
import mimetypes
from google import genai
from google.genai import types
def save_binary_file(file_name, data):
f = open(file_name, "wb")
f.write(data)
f.close()
def generate():
client = genai.Client(
api_key=os.environ.get("GEMINI_API_KEY"),
)
model = "gemini-2.0-flash-exp-image-generation"
contents = [
types.Content(
role="user",
parts=[
types.Part.from_text(text="""INSERT_INPUT_HERE"""),
],
),
]
generate_content_config = types.GenerateContentConfig(
temperature=1,
top_p=0.95,
top_k=40,
max_output_tokens=8192,
response_modalities=[
"image",
"text",
],
safety_settings=[
types.SafetySetting(
category="HARM_CATEGORY_CIVIC_INTEGRITY",
threshold="OFF", # Off
),
],
response_mime_type="text/plain",
)
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
if chunk.candidates[0].content.parts[0].inline_data:
file_name = "ENTER_FILE_NAME"
inline_data = chunk.candidates[0].content.parts[0].inline_data
file_extension = mimetypes.guess_extension(inline_data.mime_type)
save_binary_file(
f"{file_name}{file_extension}", inline_data.data
)
print(
"File of mime type"
f" {inline_data.mime_type} saved"
f"to: {file_name}"
)
else:
print(chunk.text)
if __name__ == "__main__":
generate()
|