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 def save_binary_file(file_name, data): with open(file_name, "wb") as f: f.write(data) def generate(text, file_name, api_key, model="gemini-2.0-flash-exp"): # Initialize client using provided api_key (or fallback to env variable) client = genai.Client(api_key=(api_key.strip() if api_key and api_key.strip() != "" else os.environ.get("GEMINI_API_KEY"))) try: print("Uploading file to Gemini API...") files = [ client.files.upload(file=file_name) ] contents = [ types.Content( role="user", parts=[ types.Part.from_uri( file_uri=files[0].uri, mime_type=files[0].mime_type, ), types.Part.from_text(text=text), ], ), ] generate_content_config = types.GenerateContentConfig( temperature=0, # Lower temperature for more consistent, conservative results top_p=0.92, max_output_tokens=8192, response_modalities=["image", "text"], response_mime_type="text/plain", # Additional parameters to encourage subtle, natural results safety_settings=[ { "category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_MEDIUM_AND_ABOVE" } ] ) text_response = "" image_path = None # Create a temporary file to potentially store image data with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp: temp_path = tmp.name print("Sending request to Gemini API...") # Add a timeout to prevent indefinite waiting start_time = time.time() max_wait_time = 60 # Maximum wait time in seconds try: stream = client.models.generate_content_stream( model=model, contents=contents, config=generate_content_config, ) for chunk in stream: # Check for timeout if time.time() - start_time > max_wait_time: print("Gemini API request timed out after", max_wait_time, "seconds") break if not chunk.candidates or not chunk.candidates[0].content or not chunk.candidates[0].content.parts: continue candidate = chunk.candidates[0].content.parts[0] # Check for inline image data if candidate.inline_data: save_binary_file(temp_path, candidate.inline_data.data) print(f"Smile enhancement image generated: {temp_path}") image_path = temp_path # If an image is found, we assume that is the desired output. break else: # Accumulate text response if no inline_data is present. text_response += chunk.text + "\n" print("Received text response from Gemini API") except Exception as e: print(f"Error during content generation: {str(e)}") # Continue with the function, returning empty responses except Exception as e: print(f"Error in Gemini API setup: {str(e)}") return None, f"Error: {str(e)}" finally: # Always clean up files try: if 'files' in locals() and files: del files except: pass return image_path, text_response def assess_image_quality(original_image, enhanced_image): """ Assesses the quality of the enhanced image based on specific criteria. Returns a tuple of (is_acceptable, feedback_message) """ try: # Check if enhanced image exists if enhanced_image is None: return False, "No enhanced image generated" # Image dimension checks if enhanced_image.size[0] < 100 or enhanced_image.size[1] < 100: return False, "Enhanced image appears to be too small or improperly sized" # Check that the enhanced image has similar dimensions to the original # This helps ensure facial proportions are maintained width_diff = abs(original_image.size[0] - enhanced_image.size[0]) height_diff = abs(original_image.size[1] - enhanced_image.size[1]) # If dimensions are significantly different, it suggests the image proportions changed if width_diff > 20 or height_diff > 20: return False, "Enhanced image dimensions differ significantly from original, suggesting facial proportions may have changed" # Check image has proper RGB channels for natural skin tones if enhanced_image.mode != 'RGB': return False, "Enhanced image does not have the correct color mode" # For now, we'll do basic checks and assume the model follows guidelines return True, "Image passes quality assessment criteria" except Exception as e: print(f"Error in quality assessment: {str(e)}") # Default to not accepting the image if assessment fails return False, f"Assessment error: {str(e)}" def compare_image_results(results_list): """ Compares multiple generated images and returns the best one. If no valid results, returns None. """ if not results_list or all(img is None for img in results_list): return None # Filter out None values valid_results = [img for img in results_list if img is not None] if not valid_results: return None # If there's only one valid result, return it if len(valid_results) == 1: return valid_results[0] # For now, we just return the last valid result # In a more advanced implementation, this could use computer vision techniques # to analyze facial features, smile quality, and natural appearance print(f"Comparing {len(valid_results)} valid results and selecting best one") return valid_results[-1] # Return the last attempt as potentially the best one def process_smile_enhancement(input_image, max_attempts=3): try: if input_image is None: return None, "", "" # Get API key from environment variable gemini_api_key = "AIzaSyCVzRDxkuvtaS1B22F_F-zl0ehhXR0nuU8" if not gemini_api_key: print("Error: GEMINI_API_KEY not found in environment variables") return [input_image], "", "API key not configured" # Save the input image to a temporary file with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp: input_path = tmp.name input_image.save(input_path) print(f"Input image saved to {input_path}") # Initialize attempt counter and result variables current_attempt = 0 result_images = [] # Store all generated images for comparison feedback_history = [] max_processing_time = 150 # Maximum time in seconds for overall processing start_processing_time = time.time() while current_attempt < max_attempts: # Check if overall processing time exceeded if time.time() - start_processing_time > max_processing_time: print(f"Overall processing time exceeded {max_processing_time} seconds") break current_attempt += 1 print(f"Starting processing attempt {current_attempt}/{max_attempts}...") # Create a comprehensive prompt for true smile enhancement that affects facial features naturally # Adjust prompt based on previous attempts if needed prompt = """ Create a naturally enhanced smile that focuses primarily on the overall facial expression rather than perfect teeth. Make the following personalized improvements: - Focus on enhancing the OVERALL SMILE EXPRESSION with natural eye crinkles, cheeks, and subtle facial changes - Create authentic "Duchenne smile" characteristics with proper eye corner crinkles (crow's feet) appropriate for this person's age - Enhance the natural lifting of cheeks that occurs in genuine smiles WITHOUT widening the face - Add the characteristic slight narrowing of the eyes that happens in genuine smiles - Create subtle dimples ONLY if they already exist in the original image - Boost the overall joyful expression while maintaining the person's unique facial structure - Maintain natural-looking nasolabial folds (smile lines) consistent with the smile intensity - Subtly complement existing teeth - they should remain natural looking with their original character IMPORTANT GUIDELINES: - FOCUS ON THE SMILE AS A COMPLETE FACIAL EXPRESSION - not just teeth - The most important aspects are eye crinkles, cheek raising, and natural facial expressions - Teeth should be subtly complemented but NOT the main focus of the enhancement - PRESERVE THE PERSON'S NATURAL DENTAL CHARACTERISTICS - teeth should look like THEIR teeth - Keep teeth coloration natural and appropriate for the person - avoid any artificial whitening - Maintain all natural imperfections in tooth alignment that give character to the smile - Create a genuine, authentic-looking smile that affects the entire face naturally - ABSOLUTELY CRITICAL: DO NOT widen the face or change face width/shape at all - Preserve the person's identity completely (extremely important) - Preserve exact facial proportions of the original image - Maintain natural-looking results appropriate for the person's age and face structure - Keep the original background, lighting, and image quality intact - Ensure the enhanced smile looks natural, genuine, and believable - Create a smile that looks like a moment of true happiness for THIS specific person """ # If not the first attempt, add previous feedback to the prompt if current_attempt > 1 and feedback_history: prompt += """ IMPORTANT FEEDBACK FROM PREVIOUS ATTEMPT: """ + " ".join(feedback_history) + """ Please address these issues in this new attempt. """ # Process silently print(f"Processing attempt {current_attempt}/{max_attempts}...") # Set timeout for individual API call api_call_timeout = time.time() + 45 # 45 second timeout for API call try: # Process the image using Google's Gemini model with timeout image_path, text_response = generate(text=prompt, file_name=input_path, api_key=gemini_api_key) # Check if API call timeout occurred if time.time() > api_call_timeout: print("API call timeout occurred") feedback_history.append("API call timed out, trying again with simplified request.") continue print(f"API response received: Image path: {image_path is not None}, Text length: {len(text_response)}") if image_path: # Load and convert the image if needed try: current_result = Image.open(image_path) if current_result.mode == "RGBA": current_result = current_result.convert("RGB") print("Successfully loaded generated image for attempt " + str(current_attempt)) # Assess the quality of the enhanced image is_acceptable, assessment_feedback = assess_image_quality(input_image, current_result) print(f"Image quality assessment: {is_acceptable}, {assessment_feedback}") if is_acceptable: # Store the acceptable result for later comparison result_images.append(current_result) print(f"Added acceptable result from attempt {current_attempt} to results list") # Continue with additional attempts to potentially get even better results if current_attempt < max_attempts: feedback_history.append("Previous attempt successful, trying to further improve...") continue else: # Image didn't pass quality assessment, add feedback for next attempt feedback_history.append(assessment_feedback) # Still store the result for potential use if no better options are found result_images.append(current_result) except Exception as img_error: print(f"Error processing the generated image: {str(img_error)}") feedback_history.append(f"Error with image: {str(img_error)}") else: # No image was generated, only text response print("No image was generated, only text response") feedback_history.append("No image was generated in the previous attempt.") except Exception as gen_error: print(f"Error during generation attempt {current_attempt}: {str(gen_error)}") feedback_history.append(f"Error during processing: {str(gen_error)}") # Compare all results and select the best one print(f"All attempts completed. Comparing {len(result_images)} results") if result_images: # Select the best result from all generated images best_result = compare_image_results(result_images) if best_result: print("Returning best result from multiple attempts") success_message = "Enhancement completed after multiple attempts to find the best result" return [best_result], "", success_message # Return the original image as a fallback without messages print("Returning original image as fallback - no valid results generated") return [input_image], "", "No satisfactory enhancements could be generated" except Exception as e: # Return the original image silently on error print(f"Overall error in process_smile_enhancement: {str(e)}") return [input_image], "", "" # Create a clean interface with minimal UI elements and no settings/deployment info with gr.Blocks(title="Smile Enhancement", css="footer {visibility: hidden} .gradio-container {min-height: 0 !important}") as demo: with gr.Row(): with gr.Column(): image_input = gr.Image( type="pil", label=None, image_mode="RGB", elem_classes="upload-box" ) submit_btn = gr.Button("Enhance Smile with Natural Expressions", elem_classes="generate-btn") with gr.Column(): output_gallery = gr.Gallery(label=None) # Simplify feedback to minimize UI elements feedback_text = gr.Textbox(label=None, visible=True, elem_classes="status-box") # Hidden element for structure output_text = gr.Textbox(visible=False) submit_btn.click( fn=process_smile_enhancement, inputs=[image_input], outputs=[output_gallery, output_text, feedback_text] ) # Launch the app without showing Gradio branding or share links demo.queue(max_size=50).launch( show_api=False, share=False, show_error=True, server_name="0.0.0.0", )