app_cloud_gpu_robust.py

import os
import gradio as gr
import requests
from PIL import Image
import numpy as np
import io
import json
import base64
import time
import random

# Global variables
FEATURE_TYPES = ["Eyes", "Nose", "Lips", "Face Shape", "Hair", "Body"]
MODIFICATION_PRESETS = {
    "Eyes": ["Larger", "Smaller", "Change Color", "Change Shape"],
    "Nose": ["Refine", "Reshape", "Resize"],
    "Lips": ["Fuller", "Thinner", "Change Color"],
    "Face Shape": ["Slim", "Round", "Define Jawline", "Soften Features"],
    "Hair": ["Change Color", "Change Style", "Add Volume"],
    "Body": ["Slim", "Athletic", "Curvy", "Muscular"]
}

# Mapping from our UI controls to text instructions
INSTRUCTION_MAPPING = {
    "Eyes": {
        "Larger": "make the eyes larger",
        "Smaller": "make the eyes smaller",
        "Change Color": "change the eye color to blue",
        "Change Shape": "make the eyes more almond shaped"
    },
    "Nose": {
        "Refine": "refine the nose shape",
        "Reshape": "make the nose more straight",
        "Resize": "make the nose smaller"
    },
    "Lips": {
        "Fuller": "make the lips fuller",
        "Thinner": "make the lips thinner",
        "Change Color": "make the lips more red"
    },
    "Face Shape": {
        "Slim": "make the face slimmer",
        "Round": "make the face more round",
        "Define Jawline": "define the jawline more",
        "Soften Features": "soften the facial features"
    },
    "Hair": {
        "Change Color": "change the hair color to blonde",
        "Change Style": "make the hair wavy",
        "Add Volume": "add more volume to the hair"
    },
    "Body": {
        "Slim": "make the body slimmer",
        "Athletic": "make the body more athletic",
        "Curvy": "make the body more curvy",
        "Muscular": "make the body more muscular"
    }
}

# List of available GPU Spaces for image editing
GPU_SPACES = [
    {
        "name": "InstructPix2Pix",
        "url": "https://timbrooks-instruct-pix2pix.hf.space/api/predict",
        "format_request": lambda img_str, instruction: {
            "data": [
                f"data:image/png;base64,{img_str}",  # Input image
                instruction,                         # Instruction
                50,                                  # Steps
                7.5,                                 # Text CFG
                1.5,                                 # Image CFG
                random.randint(1, 9999),            # Random Seed
                False,                               # Randomize seed
                True,                                # Fix CFG
                False                                # Randomize CFG
            ]
        },
        "parse_response": lambda response: {
            "success": response.status_code == 200,
            "data": response.json()["data"][0] if response.status_code == 200 and "data" in response.json() and len(response.json()["data"]) > 0 else None
        }
    },
    {
        "name": "SD-XL Turbo",
        "url": "https://fffiloni-sdxl-turbo.hf.space/api/predict",
        "format_request": lambda img_str, instruction: {
            "data": [
                f"data:image/png;base64,{img_str}",  # Input image
                instruction,                         # Prompt
                "",                                  # Negative prompt
                25,                                  # Steps
                1024,                                # Width
                1024,                                # Height
                1.0,                                 # Guidance scale
                0.5,                                 # Strength
                random.randint(1, 9999),            # Seed
            ]
        },
        "parse_response": lambda response: {
            "success": response.status_code == 200,
            "data": response.json()["data"][0] if response.status_code == 200 and "data" in response.json() and len(response.json()["data"]) > 0 else None
        }
    }
]

# Function to process image using cloud GPU services with fallback
def process_with_cloud_gpu(image, feature_type, modification_type, intensity, custom_prompt="", use_custom_prompt=False):
    if image is None:
        return None, "Please upload an image first."
    
    # Prepare the instruction
    if use_custom_prompt and custom_prompt:
        instruction = custom_prompt
    else:
        instruction = INSTRUCTION_MAPPING[feature_type][modification_type]
        
        # Adjust instruction based on intensity
        if intensity < 0.3:
            instruction = "slightly " + instruction
        elif intensity > 0.7:
            instruction = "dramatically " + instruction
    
    # Convert image to base64 for API request
    if isinstance(image, np.ndarray):
        image_pil = Image.fromarray(image)
    else:
        image_pil = image
        
    # Resize image if too large (most models work best with images around 512-1024px)
    width, height = image_pil.size
    max_dim = 1024
    if width > max_dim or height > max_dim:
        if width > height:
            new_width = max_dim
            new_height = int(height * (max_dim / width))
        else:
            new_height = max_dim
            new_width = int(width * (max_dim / height))
        image_pil = image_pil.resize((new_width, new_height), Image.LANCZOS)
    
    # Convert to bytes for API request
    buffered = io.BytesIO()
    image_pil.save(buffered, format="PNG")
    img_str = base64.b64encode(buffered.getvalue()).decode()
    
    # Try each GPU Space in order until one succeeds
    errors = []
    for space in GPU_SPACES:
        try:
            # Format the request according to this space's requirements
            payload = space["format_request"](img_str, instruction)
            
            # Send request with timeout
            response = requests.post(space["url"], json=payload, timeout=60)
            
            # Parse the response
            result = space["parse_response"](response)
            
            if result["success"] and result["data"]:
                # Handle base64 encoded image
                if isinstance(result["data"], str) and result["data"].startswith('data:image'):
                    # Extract the output image
                    image_data = result["data"].split(',')[1]
                    decoded_image = base64.b64decode(image_data)
                    output_image = Image.open(io.BytesIO(decoded_image))
                    return output_image, f"Edit completed successfully using {space['name']} with instruction: '{instruction}'"
            
            # If we get here, this space didn't work
            errors.append(f"{space['name']}: {response.status_code} - {response.text[:100]}...")
            
        except Exception as e:
            errors.append(f"{space['name']}: {str(e)}")
            continue
    
    # If all spaces failed, return the original image and error details
    error_msg = "All GPU Spaces failed. Details:\n" + "\n".join(errors)
    return image, error_msg

# UI Components
def create_ui():
    with gr.Blocks(title="AI-Powered Facial & Body Feature Editor") as app:
        gr.Markdown("# AI-Powered Facial & Body Feature Editor")
        gr.Markdown("Upload an image and use the controls to edit specific facial and body features using cloud GPU processing.")
        
        with gr.Row():
            with gr.Column(scale=1):
                # Input controls
                input_image = gr.Image(label="Upload Image", type="pil")
                
                with gr.Group():
                    gr.Markdown("### Feature Selection")
                    feature_type = gr.Dropdown(
                        choices=FEATURE_TYPES, 
                        label="Select Feature", 
                        value="Eyes"
                    )
                    
                    # Initialize with choices for the default feature (Eyes)
                    modification_type = gr.Dropdown(
                        choices=MODIFICATION_PRESETS["Eyes"], 
                        label="Modification Type", 
                        value="Larger"
                    )
                    
                    intensity = gr.Slider(
                        minimum=0.1, 
                        maximum=1.0, 
                        value=0.5, 
                        step=0.1, 
                        label="Intensity"
                    )
                
                with gr.Group():
                    gr.Markdown("### Custom Prompt (Advanced)")
                    use_custom_prompt = gr.Checkbox(
                        label="Use Custom Prompt", 
                        value=False
                    )
                    custom_prompt = gr.Textbox(
                        label="Custom Prompt", 
                        placeholder="e.g., make the eyes blue and add long eyelashes"
                    )
                
                edit_button = gr.Button("Apply Edit", variant="primary")
                reset_button = gr.Button("Reset")
                status_text = gr.Textbox(label="Status", interactive=False)
            
            with gr.Column(scale=1):
                # Output display
                output_image = gr.Image(label="Edited Image", type="pil")
                
                with gr.Accordion("Edit History", open=False):
                    edit_history = gr.State([])
                    history_gallery = gr.Gallery(label="Previous Edits")
                
                # Information about cloud processing
                with gr.Accordion("Cloud GPU Processing Information", open=True):
                    gr.Markdown("""
                    ### About Cloud GPU Processing
                    
                    This application uses multiple public GPU-accelerated Spaces on Hugging Face to process your images:
                    
                    1. **InstructPix2Pix** - For natural language guided image editing
                    2. **SD-XL Turbo** - For fast, high-quality image modifications
                    
                    The application will automatically try each service in order until one succeeds.
                    
                    **Benefits:**
                    - GPU-accelerated processing without local setup
                    - Automatic fallback if one service is unavailable
                    - Works on any device with internet access
                    
                    **How it works:**
                    1. Your image is sent to a GPU-accelerated Space
                    2. Your feature selections are converted to text instructions
                    3. The Space processes your image using GPU acceleration
                    4. The edited image is returned to this interface
                    
                    **Note:** Processing may take 10-30 seconds depending on server load.
                    """)
        
        # Event handlers
        def update_modification_choices(feature):
            return gr.Dropdown(choices=MODIFICATION_PRESETS[feature])
        
        feature_type.change(
            fn=update_modification_choices,
            inputs=feature_type,
            outputs=modification_type
        )
        
        edit_button.click(
            fn=process_with_cloud_gpu,
            inputs=[
                input_image, 
                feature_type, 
                modification_type, 
                intensity, 
                custom_prompt, 
                use_custom_prompt
            ],
            outputs=[output_image, status_text]
        )
        
        def reset_image():
            return None, "Image reset."
        
        reset_button.click(
            fn=reset_image,
            inputs=[],
            outputs=[output_image, status_text]
        )
        
        # Add ethical usage notice
        gr.Markdown("""
        ## Ethical Usage Notice
        
        This tool is designed for creative and personal use. Please ensure:
        
        - You have appropriate rights to edit the images you upload
        - You use this tool responsibly and respect the dignity of individuals
        - You understand that AI-generated modifications are artificial and may not represent reality
        
        By using this application, you agree to these terms.
        """)
        
    return app

# Launch the app
if __name__ == "__main__":
    app = create_ui()
    app.launch(server_name="0.0.0.0", share=False)