import gradio as gr import google.generativeai as genai import os from typing import Optional, List from huggingface_hub import whoami from PIL import Image import tempfile import io # Import io for handling in-memory binary streams # --- Google Gemini API Configuration --- # Set your Google API key as an environment variable GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY", "") if not GOOGLE_API_KEY: raise ValueError("GOOGLE_API_KEY environment variable not set.") genai.configure(api_key=GOOGLE_API_KEY) # --- Define the correct model name --- GEMINI_MODEL_NAME = 'gemini-2.5-flash-image-preview' def verify_pro_status(token: Optional[gr.OAuthToken]) -> bool: """Verifies if the user is a Hugging Face PRO user or part of an enterprise org.""" if not token: return False try: user_info = whoami(token=token.token) if user_info.get("isPro", False): return True orgs = user_info.get("orgs", []) if any(org.get("isEnterprise", False) for org in orgs): return True return False except Exception as e: print(f"Could not verify user's PRO/Enterprise status: {e}") return False # --- Backend Generation Functions --- def run_single_image_logic(prompt: str, image_path: Optional[str] = None) -> str: """Handles text-to-image or single image-to-image using Google Gemini.""" try: model = genai.GenerativeModel(GEMINI_MODEL_NAME) # Use the defined model name contents = [prompt] if image_path: input_image = Image.open(image_path) contents.append(input_image) response = model.generate_content(contents) # Access the image data correctly based on the response structure # Assuming the generated content might be in response.candidates[0].content.parts[0].inline_data.data # Or direct from response.parts if it's a single part with inline_data image_data = None if hasattr(response, 'parts') and response.parts: for part in response.parts: if hasattr(part, 'inline_data') and hasattr(part.inline_data, 'data'): image_data = part.inline_data.data break elif hasattr(response, 'candidates') and response.candidates: for candidate in response.candidates: if hasattr(candidate, 'content') and hasattr(candidate.content, 'parts') and candidate.content.parts: for part in candidate.content.parts: if hasattr(part, 'inline_data') and hasattr(part.inline_data, 'data'): image_data = part.inline_data.data break if image_data: break if not image_data: raise ValueError("No image data found in the model response.") # Save the generated image to a temporary file to return its path pil_image = Image.open(io.BytesIO(image_data)) with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmpfile: pil_image.save(tmpfile.name) return tmpfile.name except Exception as e: raise gr.Error(f"Image generation failed: {e}") def run_multi_image_logic(prompt: str, images: List[str]) -> str: """ Handles multi-image editing by sending a list of images and a prompt. """ if not images: raise gr.Error("Please upload at least one image in the 'Multiple Images' tab.") try: model = genai.GenerativeModel(GEMINI_MODEL_NAME) # Use the defined model name # The prompt should be the last part of the contents list contents = [Image.open(image_path[0]) for image_path in images] contents.append(prompt) response = model.generate_content(contents) image_data = None if hasattr(response, 'parts') and response.parts: for part in response.parts: if hasattr(part, 'inline_data') and hasattr(part.inline_data, 'data'): image_data = part.inline_data.data break elif hasattr(response, 'candidates') and response.candidates: for candidate in response.candidates: if hasattr(candidate, 'content') and hasattr(candidate.content, 'parts') and candidate.content.parts: for part in candidate.content.parts: if hasattr(part, 'inline_data') and hasattr(part.inline_data, 'data'): image_data = part.inline_data.data break if image_data: break if not image_data: raise ValueError("No image data found in the model response.") # Save the generated image to a temporary file to return its path pil_image = Image.open(io.BytesIO(image_data)) with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmpfile: pil_image.save(tmpfile.name) return tmpfile.name except Exception as e: raise gr.Error(f"Image generation failed: {e}") # --- Gradio App UI --- css = ''' #sub_title{margin-top: -35px !important} .tab-wrapper{margin-bottom: -33px !important} .tabitem{padding: 0px !important} #output{margin-top: 25px} .fillable{max-width: 980px !important} .dark .progress-text {color: white} ''' with gr.Blocks(theme=gr.themes.Citrus(), css=css) as demo: gr.HTML("

Image Generation with Google Gemini

") gr.HTML("

Hugging Face PRO users can use Google's Gemini 2.5 Flash Image Preview on this Space. Subscribe to PRO

", elem_id="sub_title") pro_message = gr.Markdown(visible=False) main_interface = gr.Column(visible=False) with main_interface: with gr.Row(): with gr.Column(scale=1): active_tab_state = gr.State(value="single") with gr.Tabs() as tabs: with gr.TabItem("Single Image", id="single") as single_tab: image_input = gr.Image( type="filepath", label="Input Image (Leave blank for text-to-image)" ) with gr.TabItem("Multiple Images", id="multiple") as multi_tab: gallery_input = gr.Gallery( label="Input Images (drop all images here)", file_types=["image"] ) prompt_input = gr.Textbox( label="Prompt", info="Tell the model what you want it to do", placeholder="A delicious looking pizza" ) generate_button = gr.Button("Generate", variant="primary") with gr.Column(scale=1): output_image = gr.Image(label="Output", interactive=False, elem_id="output") use_image_button = gr.Button("♻️ Use this Image for Next Edit") gr.Markdown("## Thank you for being a PRO! 🤗") login_button = gr.LoginButton() # --- Event Handlers --- def unified_generator( prompt: str, single_image: Optional[str], multi_images: Optional[List[str]], active_tab: str, oauth_token: Optional[gr.OAuthToken] = None, ) -> str: if not verify_pro_status(oauth_token): raise gr.Error("Access Denied. This service is for PRO users only.") if active_tab == "multiple" and multi_images: return run_multi_image_logic(prompt, multi_images) else: return run_single_image_logic(prompt, single_image) single_tab.select(lambda: "single", None, active_tab_state) multi_tab.select(lambda: "multiple", None, active_tab_state) generate_button.click( unified_generator, inputs=[prompt_input, image_input, gallery_input, active_tab_state], outputs=[output_image], ) use_image_button.click( lambda img: img, inputs=[output_image], outputs=[image_input] ) # --- Access Control Logic --- def control_access( profile: Optional[gr.OAuthProfile] = None, oauth_token: Optional[gr.OAuthToken] = None ): if not profile: return gr.update(visible=False), gr.update(visible=False) if verify_pro_status(oauth_token): return gr.update(visible=True), gr.update(visible=False) else: message = ( "## ✨ Exclusive Access for PRO Users\n\n" "Thank you for your interest! This feature is available exclusively for our Hugging Face **PRO** members.\n\n" "To unlock this and many other benefits, please consider upgrading your account.\n\n" "### [**Become a PRO Member Today!**](https://huggingface.co/pro)" ) return gr.update(visible=False), gr.update(visible=True, value=message) demo.load(control_access, inputs=None, outputs=[main_interface, pro_message]) if __name__ == "__main__": demo.queue(max_size=None, default_concurrency_limit=None) demo.launch()