app.py

import gradio as gr
import requests
import random
import json
import os
from together import Together
import base64
from PIL import Image
import io

# --- Environment Variable for Together API Key ---
TOGETHER_API_KEY = os.environ.get('TOGETHER_API_KEY')

# --- Together API Client Configuration ---
# Initialize the Together client only if the API key is available
client = Together(api_key=TOGETHER_API_KEY) if TOGETHER_API_KEY else None

# --- Pixabay API Configuration ---
PIXABAY_API_KEY = os.environ.get('PIXABAY_API_KEY')
IMAGE_API_URL = 'https://pixabay.com/api/'
VIDEO_API_URL = 'https://pixabay.com/api/videos/'
PER_PAGE = 5

# --- Function to search Pixabay ---
def search_pixabay(query: str, media_type: str, image_type: str, orientation: str, video_type: str):
    """
    Searches the Pixabay API for royalty-free stock images or videos based on user query and filters.

    Args:
        query (str): The search term for finding media. If empty, an error is returned.
        media_type (str): Specifies the type of media to search for.
                         Accepted values are "Image" or "Video".
        image_type (str): Filter results by image type (used only if media_type is "Image").
                          Accepted values: "all", "photo", "illustration", "vector".
        orientation (str): Filter results by image orientation (used only if media_type is "Image").
                           Accepted values: "all", "horizontal", "vertical".
        video_type (str): Filter results by video type (used only if media_type is "Video").
                          Accepted values: "all", "film", "animation".
    """
    if not query:
        return None, None, "Please enter a search query."

    if not PIXABAY_API_KEY:
        return None, None, "Pixabay API Key not found. Please set the PIXABAY_API_KEY environment variable."

    params = {
        'key': PIXABAY_API_KEY,
        'q': query,
        'per_page': PER_PAGE,
        'page': 1,
        'safesearch': 'true'
    }

    if media_type == "Image":
        api_url = IMAGE_API_URL
        params['image_type'] = image_type
        params['orientation'] = orientation
    elif media_type == "Video":
        api_url = VIDEO_API_URL
        params['video_type'] = video_type
    else:
        # This case should not be reachable with the Gradio Radio component
        return None, None, "Invalid media type selected."

    try:
        response = requests.get(api_url, params=params)
        response.raise_for_status()
        data = response.json()

        if data.get('totalHits', 0) == 0:
            return None, None, f"No results found for '{query}'."

        hits = data.get('hits', [])
        if not hits:
             return None, None, f"No results found for '{query}'."

        selected_hit = random.choice(hits)

        if media_type == "Image":
            image_url = selected_hit.get('largeImageURL')
            if image_url:
                # Return the image URL, None for video, and an empty status message
                return image_url, None, ""
            else:
                return None, None, "Could not retrieve large image URL."

        elif media_type == "Video":
            video_urls = selected_hit.get('videos', {})
            large_video = video_urls.get('large', {})
            video_url = large_video.get('url')

            if video_url:
                # Return None for image, the video URL, and an empty status message
                return None, video_url, ""
            else:
                 # Fallback to medium quality if large is not available
                 medium_video = video_urls.get('medium', {})
                 video_url = medium_video.get('url')
                 if video_url:
                     return None, video_url, "Using medium quality video."
                 else:
                    return None, None, "Could not retrieve video URL."

    except requests.exceptions.RequestException as e:
        return None, None, f"API request error: {e}"
    except json.JSONDecodeError:
         return None, None, "Error decoding API response."
    except Exception as e:
        # Catch any other unexpected errors
        return None, None, f"An unexpected error occurred: {e}"

# --- Together AI Image Generation Functions ---
def together_text_to_image(prompt: str):
    """
    Generates an image from text using the Together AI API.

    Args:
        prompt (str): The text prompt for image generation.

    Returns:
        str: The URL of the generated image, or an error message.
    """
    if not client:
        return "Together AI client not initialized. Please set the TOGETHER_API_KEY environment variable."
    if not prompt:
        return "Please enter a prompt for text-to-image generation."

    try:
        image_completion = client.images.generate(
            model="black-forest-labs/FLUX.1.1-pro", # Hardcoded model as requested
            width=1024,
            height=768,
            steps=40, # Hardcoded steps as requested
            prompt=prompt,
        )
        return image_completion.data[0].url
    except Exception as e:
        return f"Error generating image from text: {e}"

def together_image_to_image(image_numpy, prompt: str):
    """
    Transforms an image based on a text prompt using the Together AI API.

    Args:
        image_numpy (numpy.ndarray): The input image as a NumPy array (provided by Gradio).
        prompt (str): The text prompt for image transformation.

    Returns:
        str: The URL of the transformed image, or an error message.
    """
    if not client:
        return "Together AI client not initialized. Please set the TOGETHER_API_KEY environment variable."
    if image_numpy is None:
        return "Please upload or paste an image for image-to-image transformation."
    if not prompt:
        return "Please enter a prompt for image transformation."

    try:
        # Convert the NumPy array image to a PIL Image
        img = Image.fromarray(image_numpy.astype('uint8'), 'RGB')

        # Convert the PIL Image to base64
        buffered = io.BytesIO()
        img.save(buffered, format="PNG")
        img_base64 = base64.b64encode(buffered.getvalue()).decode("utf-8")

        image_completion = client.images.generate(
            model="black-forest-labs/FLUX.1-kontext-max", # Hardcoded model as requested
            steps=40, # Hardcoded steps as requested
            prompt=prompt,
            image_base64=img_base64
        )
        return image_completion.data[0].url
    except Exception as e:
        return f"Error transforming image: {e}"


# --- Gradio Blocks Interface Definition ---
with gr.Blocks(title="Media Generation and Search Explorer") as demo:
    gr.Markdown("## Media Generation and Search Explorer")
    gr.Markdown("Explore royalty-free media from Pixabay and generate/transform images using Together AI.")

    with gr.Tab("Pixabay Search"):
        gr.Markdown("Search for royalty-free images and videos on Pixabay.")
        gr.Warning("This requires setting the PIXABAY_API_KEY environment variable.")
        with gr.Row():
            pixabay_query_input = gr.Textbox(label="Search Query", placeholder="e.g., yellow flowers", scale=2)
            pixabay_media_type_radio = gr.Radio(["Image", "Video"], label="Media Type", value="Image", scale=1)
            pixabay_search_button = gr.Button("Search")

        with gr.Column(visible=True) as pixabay_image_options_col:
            pixabay_image_type_input = gr.Radio(["all", "photo", "illustration", "vector"], label="Image Type", value="all")
            pixabay_orientation_input = gr.Radio(["all", "horizontal", "vertical"], label="Orientation", value="all")

        with gr.Column(visible=False) as pixabay_video_options_col:
            pixabay_video_type_input = gr.Radio(["all", "film", "animation"], label="Video Type", value="all")

        pixabay_status_output = gr.Textbox(label="Status", interactive=False)

        with gr.Row():
            pixabay_image_output = gr.Image(label="Result Image (URL)", type="filepath", interactive=False)
            pixabay_video_output = gr.Video(label="Result Video (URL)", type="filepath", interactive=False)

        # Logic to toggle visibility of input columns and output components based on media type selection
        def update_pixabay_inputs_blocks(media_type):
            if media_type == "Image":
                return gr.update(visible=True), gr.update(visible=False), gr.update(visible=True), gr.update(visible=False)
            elif media_type == "Video":
                return gr.update(visible=False), gr.update(visible=True), gr.update(visible=False), gr.update(visible=True)
            else:
                return gr.update(visible=True), gr.update(visible=False), gr.update(visible=True), gr.update(visible=False)

        # Trigger the update_pixabay_inputs_blocks function when media_type_radio changes
        pixabay_media_type_radio.change(
            fn=update_pixabay_inputs_blocks,
            inputs=pixabay_media_type_radio,
            outputs=[pixabay_image_options_col, pixabay_video_options_col, pixabay_image_output, pixabay_video_output]
        )

        # Trigger the search_pixabay function when the search button is clicked
        pixabay_search_button.click(
            fn=search_pixabay,
            inputs=[
                pixabay_query_input,
                pixabay_media_type_radio,
                pixabay_image_type_input,
                pixabay_orientation_input,
                pixabay_video_type_input
            ],
            outputs=[pixabay_image_output, pixabay_video_output, pixabay_status_output]
        )

    with gr.Tab("Together AI - Text to Image"):
        gr.Markdown("Generate an image from a text prompt using Together AI.")
        gr.Warning("This requires setting the TOGETHER_API_KEY environment variable.")
        with gr.Row():
            together_text_to_image_prompt = gr.Textbox(label="Enter your prompt", scale=2)
            together_text_to_image_button = gr.Button("Generate Image", scale=1)
        together_text_to_image_output = gr.Image(label="Generated Image (URL)", type="filepath", interactive=False)

        together_text_to_image_button.click(
            fn=together_text_to_image,
            inputs=together_text_to_image_prompt,
            outputs=together_text_to_image_output,
        )

    with gr.Tab("Together AI - Image to Image"):
        gr.Markdown("Transform an uploaded image based on a text prompt using Together AI.")
        gr.Warning("This requires setting the TOGETHER_API_KEY environment variable.")
        with gr.Row():
            together_image_input = gr.Image(label="Upload or paste an image", scale=2)
            together_image_to_image_prompt = gr.Textbox(label="Enter your transformation prompt", scale=2)
            together_image_to_image_button = gr.Button("Transform Image", scale=1)
        together_image_to_image_output = gr.Image(label="Transformed Image (URL)", type="filepath", interactive=False)

        together_image_to_image_button.click(
            fn=together_image_to_image,
            inputs=[together_image_input, together_image_to_image_prompt],
            outputs=together_image_to_image_output,
        )

# --- Launch the Gradio app ---
if __name__ == "__main__":
    demo.launch(mcp_server=True)