scmlewis/image_edit_generation

1import os2import tempfile3from PIL import Image4import gradio as gr5from google import genai6from google.genai import types7 8# Helpers9def save_binary_file(file_name, data):10    with open(file_name, "wb") as f:11        f.write(data)12 13def generate_edit(prompt, pil_image, api_key, model="gemini-2.0-flash-exp"):14    client = genai.Client(api_key=(api_key.strip() if api_key and api_key.strip() != "" else os.environ.get("GEMINI_API_KEY")))15    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_img:16        image_path = tmp_img.name17        pil_image.save(image_path)18 19    files = [client.files.upload(file=image_path)]20    contents = [21        types.Content(22            role="user",23            parts=[24                types.Part.from_uri(file_uri=files[0].uri, mime_type=files[0].mime_type),25                types.Part.from_text(text=prompt),26            ],27        ),28    ]29    generate_content_config = types.GenerateContentConfig(30        temperature=1,31        top_p=0.95,32        top_k=40,33        max_output_tokens=8192,34        response_modalities=["image", "text"],35        response_mime_type="text/plain",36    )37 38    text_response = ""39    image_out_path = None40 41    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_out:42        out_path = tmp_out.name43        for chunk in client.models.generate_content_stream(44            model=model,45            contents=contents,46            config=generate_content_config,47        ):48            if not chunk.candidates or not chunk.candidates[0].content or not chunk.candidates[0].content.parts:49                continue50            candidate = chunk.candidates[0].content.parts[0]51            if candidate.inline_data:52                save_binary_file(out_path, candidate.inline_data.data)53                image_out_path = out_path54                break55            else:56                text_response += chunk.text + "\n"57    del files58    return image_out_path, text_response59 60def process_image_and_prompt(pil_image, prompt, api_key):61    try:62        image_path, text_out = generate_edit(prompt, pil_image, api_key)63        if image_path:64            img = Image.open(image_path)65            if img.mode == "RGBA":66                img = img.convert("RGB")67            return img # Return only the image on success68        else:69            # If no image generated, raise an error for Gradio popup70            raise gr.Error(f"⚠️ Image generation failed: {text_out.strip() if text_out.strip() else 'No specific error message.'}")71    except Exception as e:72        # Catch any other exceptions and re-raise as Gradio error73        raise gr.Error(f"❌ Generation failed: {str(e)}")74 75def reset_inputs(api_key_value=None):76    # Reset all inputs, keeping API key unchanged77    return None, "", api_key_value or ""78 79# Styles with gradient background for body/app container80css_style = """81:root { --bg: #0f111a; --panel: #1b1e28; --text: #e9eefc; --muted: #9fb3c8; --accent: #6a8efd; }82body, .app-container {83  /* Gradient background instead of solid */84  background: linear-gradient(135deg, #2a3a67, #1b254b);85  color: var(--text);86}87.header-block { width: 100%; display: flex; justify-content: center; padding: 8px 0; }88.header-gradient { width: 100%; padding: 20px 0; border-radius: 12px; background: linear-gradient(90deg, #6a8efd, #44abc7); text-align: center; }89.header-title { margin: 0; font-size: 2.6rem; font-weight: 900; color: #fff; text-shadow: 0 2px 8px rgba(0,0,0,.25); }90.header-subtitle { margin-top: 6px; font-size: 1.05rem; color: #e8f0ff; }91 92.main { display: flex; gap: 20px; align-items: stretch; padding: 0 12px; }93.sidebar { width: 320px; background: #1a1e2a; padding: 14px; border-radius: 12px; min-height: 360px; box-shadow: 0 2px 10px rgb(0 0 0 / 0.25); }94.sidebar h2 { color: #8ab4ff; font-size: 1rem; margin: 6px 0; }95.sidebar ul { margin: 0; padding-left: 18px; color: #d6e3ff; line-height: 1.9; }96.sidebar a { color: #97b7ff; text-decoration: none; }97.sidebar a:hover { text-decoration: underline; }98 99.main-panel { flex: 1; display: flex; flex-direction: column; gap: 12px; }100 101.section-header { font-weight: 800; font-size: 1.04rem; color: #cbd5e1; margin: 6px 0; }102 103.layout-two-col {104  display: grid;105  grid-template-columns: 1fr 1fr;106  gap: 14px;107}108@media (max-width: 1100px) {109  .layout-two-col { grid-template-columns: 1fr; }110}111 112#output-viewport { display: flex; justify-content: center; align-items: center; min-height: 260px; }113#output-image { display: flex; justify-content: center; align-items: center; }114#output-image img { max-width: 100%; max-height: 420px; object-fit: contain; border-radius: 12px; background: #23252b; }115"""116 117# Layout118with gr.Blocks(css=css_style) as app:119    gr.HTML("""120    <div class='header-block'>121      <div class='header-gradient'>122        <h1 class='header-title'>🖼️ Image Editor <span style="font-size:1.1em;">(Powered by Gemini)</span> 🔮</h1>123        <div class='header-subtitle'>Step-by-step prompts for image editing</div>124      </div>125    </div>126    """)127 128    with gr.Row():129        with gr.Column(scale=3, elem_classes="sidebar"):130            gr.Markdown(131                """132                <h2>📖 How to Use</h2>133                <ul>134                  <li>Step-by-step prompts guide the editing process.</li>135                  <li>Upload a PNG image, enter a prompt, then generate.</li>136                  <li>Keep your Gemini API key secure.</li>137                </ul>138                <hr>139                <h2>🔑 API Key</h2>140                <div>Get your key here: <a href="https://aistudio.google.com/apikey" target="_blank">Get your Google API key</a></div>141                """142            )143        with gr.Column(scale=9, elem_classes="main-panel"):144            # Step 1 & Step 3 side-by-side145            with gr.Row():146                with gr.Column():147                    gr.Markdown("<div class='section-header'>Step 1: Upload Image</div>")148                    image_input = gr.Image(type="pil", label=None, image_mode="RGBA")149                with gr.Column():150                    gr.Markdown("<div class='section-header'>Step 3: Image Output</div>")151                    output_image = gr.Image(label=None, show_label=False, type="pil")152            # Step 2: Prompt + API153            gr.Markdown("<div class='section-header'>Step 2: Enter Editing Prompt</div>")154            prompt_input = gr.Textbox(label="Edit Prompt", placeholder="Describe how to edit the image", lines=2)155            api_key_input = gr.Textbox(label="Gemini API Key (required)", placeholder="Enter your Gemini API key here", type="password")156 157            with gr.Row():158                submit_btn = gr.Button("Generate Edit", elem_classes="gradient-button")159                reset_btn = gr.Button("Reset Inputs")160            161            # Note: Status bar elements removed as requested. Errors will now show as Gradio popups.162 163            def on_submit(pil_img, prompt, key):164                if not key or key.strip() == "":165                    raise gr.Error("Gemini API Key is required!")166                # process_image_and_prompt now raises gr.Error directly for failures167                return process_image_and_prompt(pil_img, prompt, key)168 169            submit_btn.click(170                fn=on_submit,171                inputs=[image_input, prompt_input, api_key_input],172                outputs=[output_image] # Only output the image173            )174 175            reset_btn.click(176                fn=reset_inputs,177                inputs=[api_key_input],178                outputs=[image_input, prompt_input, api_key_input] # Remove status_bar from outputs179            )180 181app.launch()182