Spaces:
Sleeping
Sleeping
File size: 2,781 Bytes
8e6ca2b e5c238d 8e6ca2b e5c238d 8e6ca2b 68780eb 8e6ca2b 68780eb 8e6ca2b 68780eb e5c238d 8e6ca2b 68780eb d2236a2 8e6ca2b 68780eb 8e6ca2b 68780eb 8e6ca2b 49012e3 |
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 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 |
from flask import Flask, render_template, request, send_file
import os
import mimetypes
from google import genai
from google.genai import types
from io import BytesIO
app = Flask(__name__)
# Initialize Gemini client
client = genai.Client(api_key=GEMINI_API_KEY)
def save_binary_file(file_name, data):
"""Save binary data to a file."""
with open(file_name, "wb") as f:
f.write(data)
def generate_gemini_output(user_input, image):
model = "gemini-2.0-flash-exp-image-generation"
# Upload the image to Gemini
files = []
if image:
uploaded_file = client.files.upload(file=image)
files.append(uploaded_file)
# Prepare the input content
parts = []
if files:
parts.append(types.Part.from_uri(file_uri=files[0].uri, mime_type=files[0].mime_type))
# Incorporate Gemini magic: modify the prompt to instruct removal of the user-specified element.
if user_input:
magic_prompt = f"Remove {user_input} from the image"
parts.append(types.Part.from_text(text=magic_prompt))
contents = [types.Content(role="user", parts=parts)]
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"),
],
response_mime_type="text/plain",
)
result_text = None
result_image = None
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
part = chunk.candidates[0].content.parts[0]
if part.inline_data:
file_name = "generated_output"
file_extension = mimetypes.guess_extension(part.inline_data.mime_type) or ".png"
file_path = os.path.join("static", f"{file_name}{file_extension}")
save_binary_file(file_path, part.inline_data.data)
result_image = file_path
else:
result_text = part.text
return result_text, result_image
@app.route("/", methods=["GET", "POST"])
def index():
result_text = None
result_image = None
if request.method == "POST":
user_input = request.form.get("user_input")
image = request.files.get("image_input")
result_text, result_image = generate_gemini_output(user_input, image)
return render_template("index.html", result_text=result_text, result_image=result_image)
if __name__ == "__main__":
app.run(host="0.0.0.0", port=7860) |