app.py

'''import requests
from io import BytesIO
from PIL import Image
import gradio as gr
import os

# Hugging Face Inference API setup
API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-2"
API_TOKEN = os.getenv("HF_TOKEN")  # Loaded from Spaces secrets
HEADERS = {"Authorization": f"Bearer {API_TOKEN}"}

def generate_image(prompt):
    """Generate an image from a text prompt using Hugging Face API."""
    if not API_TOKEN:
        return None, "Error: HF_API_TOKEN not set in Spaces secrets!"
    
    if not prompt or not prompt.strip():
        return None, "Error: Please enter a valid prompt!"
    
    payload = {
        "inputs": prompt,
        "options": {"wait_for_model": True}
    }
    
    try:
        response = requests.post(API_URL, headers=HEADERS, json=payload, timeout=30)
        response.raise_for_status()
        image_bytes = response.content
        image = Image.open(BytesIO(image_bytes))
        return image, "Image generated successfully!"
    except requests.exceptions.Timeout:
        return None, "Error: API request timed out after 30 seconds. Try again later."
    except requests.exceptions.RequestException as e:
        return None, f"API request failed: {str(e)}"
    except Exception as e:
        return None, f"Image processing failed: {str(e)}"

# Gradio interface
with gr.Blocks(title="Instant Image Generator") as demo:
    gr.Markdown("# Instant Image Generator")
    gr.Markdown("Enter a text prompt to generate an image in seconds!")
    
    with gr.Row():
        prompt_input = gr.Textbox(label="Prompt", placeholder="A futuristic city at sunset")
        submit_btn = gr.Button("Generate Image")
    
    output_image = gr.Image(label="Generated Image")
    status = gr.Textbox(label="Status", interactive=False)
    
    # Connect button to function
    submit_btn.click(
        fn=generate_image,
        inputs=prompt_input,
        outputs=[output_image, status]
    )
    
    # Diagnostic check
    if not API_TOKEN:
        gr.Warning("API token missing! Set HF_API_TOKEN in Spaces Settings > Variables and Secrets.")

demo.launch()'''





       
import requests
from io import BytesIO
import base64
from PIL import Image
import gradio as gr
import os
import numpy as np

# Hugging Face Inference API setup
API_URL_TEXT = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-2"  # For text-to-image
API_URL_INPAINT = "https://api-inference.huggingface.co/models/runwayml/stable-diffusion-inpainting"  # For image-to-image with face
API_TOKEN = os.getenv("HF_TOKEN")  # Loaded from environment variable or Spaces secrets
HEADERS = {"Authorization": f"Bearer {API_TOKEN}" if API_TOKEN else ""}

def image_to_base64(image):
    """Convert a PIL Image to a base64-encoded string."""
    if image is None:
        return None
    buffered = BytesIO()
    image.save(buffered, format="PNG")
    return base64.b64encode(buffered.getvalue()).decode("utf-8")

def base64_to_image(base64_string):
    """Convert a base64-encoded string to a PIL Image."""
    if not base64_string or base64_string.strip() == "":
        return None
    try:
        img_data = base64.b64decode(base64_string)
        return Image.open(BytesIO(img_data))
    except Exception as e:
        return None

def generate_image_with_face(prompt, input_image):
    """Generate an image from a text prompt using an optional uploaded face image."""
    # Check API token
    if not API_TOKEN:
        return None, "Error: HF_TOKEN not set in environment variables or Spaces secrets! Please set it and retry."
    
    # Validate prompt
    if not prompt or not prompt.strip():
        return None, "Error: Please enter a valid prompt!"

    # Text-to-image path (no face image uploaded)
    if input_image is None or (isinstance(input_image, np.ndarray) and input_image.size == 0):
        payload = {
            "inputs": prompt,
            "options": {"wait_for_model": True}
        }
        
        try:
            response = requests.post(
                API_URL_TEXT,
                headers=HEADERS,
                json=payload,
                timeout=120  # Increased timeout to 120 seconds for text-to-image
            )
            response.raise_for_status()
            
            # Handle the response (could be bytes or JSON with base64)
            if response.headers.get("content-type") == "application/json":
                data = response.json()
                if isinstance(data, str):  # If the response is a base64 string
                    image = base64_to_image(data)
                    if image:
                        return image, "Image generated successfully using text prompt!"
                    return None, "Error: Failed to decode API response as an image."
                elif isinstance(data, dict) and "image" in data:
                    image = base64_to_image(data["image"])
                    if image:
                        return image, "Image generated successfully using text prompt!"
                    return None, "Error: Failed to decode API response as an image."
                return None, "Error: Unexpected API response format."
            
            # If response is bytes (common for image data)
            image_bytes = response.content
            image = Image.open(BytesIO(image_bytes))
            return image, "Image generated successfully using text prompt!"
        
        except requests.exceptions.Timeout:
            return None, "Error: API request timed out after 120 seconds. Try again later or check your internet connection."
        except requests.exceptions.HTTPError as e:
            return None, f"API error: HTTP {e.response.status_code} - {e.response.text}"
        except requests.exceptions.RequestException as e:
            return None, f"API request failed: {str(e)}"
        except Exception as e:
            return None, f"Image processing failed: {str(e)}"
    
    # Image-to-image (inpainting) path with face upload
    try:
        # Convert input image to PIL and resize to 512x512 (Stable Diffusion's default)
        input_pil = Image.fromarray(input_image).convert("RGB").resize((512, 512))
        
        # Create a full white mask for inpainting (simplified for now)
        mask = Image.new("L", (512, 512), 255)  # White mask for full inpainting
        
        # Convert images to base64 for JSON serialization
        image_base64 = image_to_base64(input_pil)
        mask_base64 = image_to_base64(mask)
        
        if image_base64 is None or mask_base64 is None:
            return None, "Error: Failed to process the uploaded image or mask."

        # Prepare payload for inpainting
        payload = {
            "inputs": {
                "image": image_base64,  # Base64-encoded face image
                "mask": mask_base64,    # Base64-encoded mask
            },
            "parameters": {
                "prompt": prompt,
                "negative_prompt": "blurry, distorted, low quality, extra people, dark lighting",
                "strength": 0.8,  # Balance between input preservation and generation
                "num_inference_steps": 50  # Quality vs. speed tradeoff
            },
            "options": {"wait_for_model": True}
        }
        
        response = requests.post(
            API_URL_INPAINT,
            headers=HEADERS,
            json=payload,
            timeout=120  # Keep timeout at 120 seconds for image-to-image
        )
        response.raise_for_status()
        
        # Handle the response (could be bytes or JSON with base64)
        if response.headers.get("content-type") == "application/json":
            data = response.json()
            if isinstance(data, str):  # If the response is a base64 string
                image = base64_to_image(data)
                if image:
                    return image, "Image generated successfully using your face and prompt!"
                return None, "Error: Failed to decode API response as an image."
            elif isinstance(data, dict) and "image" in data:
                image = base64_to_image(data["image"])
                if image:
                    return image, "Image generated successfully using your face and prompt!"
                return None, "Error: Failed to decode API response as an image."
            return None, "Error: Unexpected API response format."
        
        # If response is bytes (common for image data)
        image_bytes = response.content
        image = Image.open(BytesIO(image_bytes))
        return image, "Image generated successfully using your face and prompt!"
    
    except requests.exceptions.Timeout:
        return None, "Error: API request timed out for image processing after 120 seconds. Try again later or simplify the prompt."
    except requests.exceptions.HTTPError as e:
        return None, f"API error: HTTP {e.response.status_code} - {e.response.text}"
    except requests.exceptions.RequestException as e:
        return None, f"API request failed: {str(e)}"
    except Exception as e:
        return None, f"Image processing failed: {str(e)}"

# Gradio interface
with gr.Blocks(title="Face-Based Image Generator") as demo:
    gr.Markdown("# Face-Based Image Generator")
    gr.Markdown("Enter a text prompt and optionally upload your face image to generate a custom image!")
    
    with gr.Row():
        prompt_input = gr.Textbox(label="Prompt", placeholder="A joyful person at a birthday party")
        image_input = gr.Image(label="Upload Your Face Image (Optional)", type="numpy")
        submit_btn = gr.Button("Generate Image")
    
    output_image = gr.Image(label="Generated Image")
    status = gr.Textbox(label="Status", interactive=False)
    
    # Connect button to function
    submit_btn.click(
        fn=generate_image_with_face,
        inputs=[prompt_input, image_input],
        outputs=[output_image, status]
    )
    
    # Diagnostic check
    if not API_TOKEN:
        gr.Warning("API token missing! Set HF_TOKEN as an environment variable (e.g., in .env or terminal) or in Spaces secrets.")

demo.launch()