import json import os import time import uuid import tempfile from PIL import Image, ImageDraw, ImageFont import gradio as gr import base64 import mimetypes from google import genai from google.genai import types # Helper function to save binary data def save_binary_file(file_name, data): with open(file_name, "wb") as f: f.write(data) # Function to get tags from an image using Gemini def get_image_tags(file_name, text_prompt, model="gemini-2.0-flash-exp"): """ Analyzes an image using a text prompt and returns the text response. Used specifically for generating tags in this case. """ api_key = os.environ.get("geminigoogle") if not api_key: # Return a clear message if API key is missing return "Error: GEMINI_API_KEY environment variable (geminigoogle) not set for tagging." client = genai.Client(api_key=api_key) uploaded_files = [] # Keep track of uploaded files for cleanup try: # Upload the file uploaded_files = [client.files.upload(file=file_name)] print(f"Uploaded file for tagging: {uploaded_files[0].uri}") contents = [ types.Content( role="user", parts=[ types.Part.from_uri( file_uri=uploaded_files[0].uri, mime_type=uploaded_files[0].mime_type, ), types.Part.from_text(text=text_prompt), ], ), ] # Configure for text-only response (focus on getting JSON) generate_content_config = types.GenerateContentConfig( temperature=0.5, # Lower temperature might give more focused tags top_p=0.95, top_k=40, max_output_tokens=1024, response_modalities=["text"], response_mime_type="text/plain", # Expect plain text ) # Use generate_content for a single text response response = client.models.generate_content( model=model, contents=contents, config=generate_content_config, ) tag_response = "" if response and response.candidates and response.candidates[0].content and response.candidates[0].content.parts: # Concatenate all text parts from the response for part in response.candidates[0].content.parts: if hasattr(part, 'text'): tag_response += part.text else: tag_response = "Could not generate tags." return tag_response except Exception as e: print(f"Error during tagging API call: {e}") return f"Error generating tags: {e}" finally: for file in uploaded_files: try: client.files.delete(name=file.name) print(f"Deleted uploaded file after tagging: {file.name}") except Exception as cleanup_e: print(f"Error deleting uploaded file {file.name}: {cleanup_e}") # Function for the main image processing call def generate(text, file_name, model="gemini-2.0-flash-exp"): api_key = os.environ.get("geminigoogle") if not api_key: raise ValueError("GEMINI_API_KEY environment variable (geminigoogle) not set.") client = genai.Client(api_key=api_key) uploaded_files = [] temp_output_image_path = None try: uploaded_files = [client.files.upload(file=file_name)] print(f"Uploaded file for generation: {uploaded_files[0].uri}") contents = [ types.Content( role="user", parts=[ types.Part.from_uri( file_uri=uploaded_files[0].uri, mime_type=uploaded_files[0].mime_type, ), types.Part.from_text(text=text), ], ), ] generate_content_config = types.GenerateContentConfig( temperature=1, top_p=0.95, top_k=40, max_output_tokens=8192, response_modalities=["image", "text"], # Expecting potentially image and text response_mime_type="text/plain", ) text_response = "" image_path = None with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp: temp_output_image_path = tmp.name print("Starting generation stream...") # Stream the response 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 for part in chunk.candidates[0].content.parts: # Check for text parts text_part = getattr(part, "text", "") if text_part: text_response += text_part if part.inline_data: print(f"Received image data with mime type {part.inline_data.mime_type}. Saving to {temp_output_image_path}") save_binary_file(temp_output_image_path, part.inline_data.data) image_path = temp_output_image_path # Set the output image path print("Generation stream finished.") if not image_path or not os.path.exists(image_path) or os.path.getsize(image_path) == 0: print("No valid image data was received or saved.") image_path = None return image_path, text_response.strip() except Exception as e: print(f"Error during main generation API call: {e}") if temp_output_image_path and os.path.exists(temp_output_image_path): os.remove(temp_output_image_path) raise e # Re-raise the exception after cleanup finally: for file in uploaded_files: try: client.files.delete(name=file.name) print(f"Deleted uploaded file after generation: {file.name}") except Exception as cleanup_e: print(f"Error deleting uploaded file {file.name}: {cleanup_e}") # Main processing function for Gradio def process_image_and_prompt(composite_pil, prompt, enable_tagging=True): composite_path = None # Initialize input temp file path for finally block output_image_path = None # Initialize output temp file path for finally block try: # 1. Save the input PIL image to a temporary file with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp: composite_path = tmp.name if composite_pil.mode == "RGBA": composite_pil.save(composite_path, format="PNG") else: composite_pil.save(composite_path, format="PNG") # Save as PNG by default file_name = composite_path # This is the path to the saved input image file model = "gemini-2.0-flash-exp" # Specify the model here # 2. Call get_image_tags to get tags from the original image tag_json_string = "" if enable_tagging: tagging_prompt = "Analyze this image. Provide a JSON object containing a single key, 'tags', whose value is a JSON array of strings, representing relevant keywords or tags for the image content. Example: {\"tags\": [\"apple\", \"fruit\", \"red\"]}. Provide ONLY the JSON object and nothing else." tag_json_string = get_image_tags(file_name, tagging_prompt, model=model) # 3. Call generate for the main image processing based on the user prompt output_image_path, main_text_response = generate(text=prompt, file_name=file_name, model=model) # 4. Combine the tag JSON string and the main text response # Format the output clearly if tag_json_string: final_text_output = f"{tag_json_string},{main_text_response}" else: final_text_output = main_text_response # 5. Prepare the image output for the Gradio gallery result_img = None image_output_list = [] if output_image_path and os.path.exists(output_image_path): try: result_img = Image.open(output_image_path) # Convert to RGB for display if it's RGBA (Gradio Gallery often expects RGB) if result_img.mode == "RGBA": result_img = result_img.convert("RGB") image_output_list = [result_img] # Add the image to the list for the gallery except Exception as img_e: print(f"Error opening generated image {output_image_path}: {img_e}") # If image opening fails, don't return an image image_output_list = [] # Append error to text response final_text_output += f"\n\n---\n\nError loading generated image: {img_e}" # 6. Return results to Gradio return image_output_list, final_text_output except Exception as e: # Catch any exceptions during the process print(f"An error occurred during processing: {e}") # Use Gradio's error handling to display a message in the UI raise gr.Error(f"Processing failed: {e}", duration=5) finally: # 7. Clean up temporary files # Clean up the temporary input file if composite_path and os.path.exists(composite_path): try: os.remove(composite_path) print(f"Deleted input temporary file: {composite_path}") except Exception as cleanup_e: print(f"Error deleting input temporary file {composite_path}: {cleanup_e}") if output_image_path and os.path.exists(output_image_path): try: os.remove(output_image_path) print(f"Deleted output temporary file: {output_image_path}") except Exception as cleanup_e: print(f"Error deleting output temporary file {output_image_path}: {cleanup_e}") # Gradio інтерфейс (unchanged from your original code, except connection) with gr.Blocks( # css_paths="style.css", # Тимчасово закоментували цей рядок ) as demo: gr.HTML( """
Gemini logo

Gemini for Image Editing

Powered by Gradio⚡️| Duplicate this Repo | Get an API Key | Follow me on Twitter: Ameerazam18

""" ) # Прибираємо секцію API Configuration або змінюємо її опис, оскільки ключ більше не вводиться with gr.Accordion("⚠️ API Configuration ⚠️", open=False, elem_classes="config-accordion"): gr.Markdown(""" - **Ваш Gemini API ключ має бути збережений у змінній оточення `geminigoogle` в налаштуваннях Hugging Face Space.** - ❗ Іноді модель повертає текст замість зображення. """) with gr.Accordion("📌 Usage Instructions", open=False, elem_classes="instructions-accordion"): gr.Markdown(""" ### 📌 Usage - Upload an image and enter a prompt to generate outputs. - The response will include generated tags for the original image (in JSON format) and Gemini's text output. - If an edited image is returned, it will appear in the gallery. If not, only text will appear. - Upload Only PNG Image - ❌ **Do not use NSFW images!** """) with gr.Row(elem_classes="main-content"): with gr.Column(elem_classes="input-column"): image_input = gr.Image( type="pil", label="Upload Image", image_mode="RGBA", elem_id="image-input", elem_classes="upload-box" ) prompt_input = gr.Textbox( lines=2, placeholder="Enter prompt here (e.g., 'change text to \"HELLO\"', 'remove the background')", label="Prompt for Gemini", elem_classes="prompt-input" ) with_tags = gr.Checkbox(label="Enable Tagging", value=True) submit_btn = gr.Button("Generate", elem_classes="generate-btn") with gr.Column(elem_classes="output-column"): output_gallery = gr.Gallery(label="Generated Image Output", elem_classes="output-gallery", allow_preview=True) output_text = gr.Textbox( label="Gemini Output (Tags + Response)", placeholder="Original image tags (JSON) and Gemini's response will appear here.", elem_classes="output-text", lines=10 # Give more space for combined output ) # Connect the button click to the updated processing function submit_btn.click( fn=process_image_and_prompt, inputs=[image_input, prompt_input, with_tags], outputs=[output_gallery, output_text], ) gr.Markdown("## Try these examples", elem_classes="gr-examples-header") examples = [ ["data/1.webp", 'change text to "AMEER"'], ["data/2.webp", "remove the spoon from hand only"], ["data/3.webp", 'change text to "Make it "'], ["data/1.jpg", "add joker style only on face"], ["data/1777043.jpg", "add joker style only on face"], ["data/2807615.jpg", "add lipstick on lip only"], ["data/76860.jpg", "add lipstick on lip only"], ["data/2807615.jpg", "make it happy looking face only"], ] gr.Examples( examples=examples, inputs=[image_input, prompt_input], elem_id="examples-grid" ) demo.queue(max_size=50).launch()