app.py

import gradio as gr
import json
import os
import mimetypes
from google import genai
from google.genai import types
from PIL import Image
import time

# --- Helper Function to Save Generated Image ---
def save_binary_file(directory, file_name, data):
    """Saves binary data to a file, creating the directory if needed."""
    if not os.path.exists(directory):
        os.makedirs(directory)
    file_path = os.path.join(directory, file_name)
    with open(file_path, "wb") as f:
        f.write(data)
    print(f"File saved to: {file_path}")
    return file_path

# --- Main Function to Generate Image ---
def generate_image(
    api_key,
    reference_image,
    scene,
    subject_type,
    age_range,
    hair,
    makeup,
    jewellery,
    top,
    bottom,
    footwear,
    wardrobe_notes,
    pose_angle,
    body_pose,
    hands_pose,
    framing,
    camera_device,
    flash,
    orientation,
    aspect_ratio,
    distance,
    focus,
    texture,
    sharpness,
    color,
    effects,
    background_environment,
    background_props,
    style_genre,
    authenticity,
    use_original_structure,
    face_description,
    ban_mirror,
    ban_phone,
    ban_selfie,
    ban_grainy,
    ban_harsh_flash,
    ban_logos,
    ban_nsfw,
    ban_cropped_feet,
    output_count,
    output_size,
    safety,
    variant_name,
    variant_angle,
):
    # --- Input Validation ---
    if not api_key:
        raise gr.Error("API Key is missing. Please enter your Gemini API key.")
    if reference_image is None:
        raise gr.Error("Reference image is missing. Please upload an image.")

    # --- Build Banned List ---
    banned_items = []
    if ban_mirror: banned_items.append("mirror")
    if ban_phone: banned_items.append("phone")
    if ban_selfie: banned_items.append("selfie look")
    if ban_grainy: banned_items.append("grainy noise")
    if ban_harsh_flash: banned_items.append("harsh LED flash")
    if ban_logos: banned_items.append("logos/brand text")
    if ban_nsfw: banned_items.append("nsfw")
    if ban_cropped_feet: banned_items.append("cropped feet")

    # --- Construct JSON Payload ---
    output_json = {
        "scene": scene,
        "subject": {"type": subject_type, "age_range": age_range, "hair": hair, "makeup": makeup, "jewellery": jewellery},
        "wardrobe": {"top": top, "bottom": bottom, "footwear": footwear, "notes": wardrobe_notes},
        "pose": {"angle": pose_angle, "body": body_pose, "hands": hands_pose, "framing": framing},
        "camera": {"device": camera_device, "flash": flash, "orientation": orientation, "aspect_ratio": aspect_ratio, "distance": distance, "focus": focus},
        "look": {"texture": texture, "sharpness": sharpness, "color": color, "effects": effects},
        "background": {"environment": background_environment, "props": background_props},
        "style": {"genre": style_genre, "authenticity": authenticity},
        "reference_face": {"use_original_structure": use_original_structure, "description": face_description},
        "ban": banned_items,
        "output": {"count": int(output_count), "size": output_size, "safety": safety},
        "variants": [{"name": variant_name, "angle": variant_angle}],
    }
    final_json_string = json.dumps(output_json, indent=4)

    # --- Call Gemini API ---
    try:
        # Configure the client
        client = genai.Client(api_key=api_key)

        # Prepare the prompt parts (JSON instructions + reference image)
        prompt_text_part = types.Part.from_text(text=final_json_string)
        
        with open(reference_image, 'rb') as f:
            image_data = f.read()
        image_mime_type = mimetypes.guess_type(reference_image)[0]
        image_part = types.Part.from_data(data=image_data, mime_type=image_mime_type)

        # Define the model and generation config
        model = "gemini-1.5-flash-latest" # Using a standard available model name
        contents = [types.Content(role="user", parts=[prompt_text_part, image_part])]
        generate_content_config = types.GenerateContentConfig(
            response_modalities=["IMAGE", "TEXT"],
        )

        # --- Process Streaming Response ---
        output_files = []
        output_directory = "generated_images"
        timestamp = int(time.time())
        file_index = 0
        
        # Make the streaming API call
        response_stream = client.models.generate_content_stream(
            model=model,
            contents=contents,
            config=generate_content_config,
        )

        for chunk in response_stream:
            if chunk.candidates and chunk.candidates[0].content and chunk.candidates[0].content.parts:
                part = chunk.candidates[0].content.parts[0]
                if part.inline_data and part.inline_data.data:
                    inline_data = part.inline_data
                    file_extension = mimetypes.guess_extension(inline_data.mime_type)
                    file_name = f"output_{timestamp}_{file_index}{file_extension}"
                    
                    # Save the file and get its path
                    saved_file_path = save_binary_file(output_directory, file_name, inline_data.data)
                    output_files.append(saved_file_path)
                    
                    file_index += 1
                elif part.text:
                    print(f"Received text chunk: {part.text}")


        if not output_files:
            return None, final_json_string, "No image was generated. Please check the model's response or your prompt."

        # Return file paths for the Gallery and the JSON for inspection
        return output_files, final_json_string, "Image generation complete."

    except Exception as e:
        # Handle potential errors gracefully
        error_message = f"An error occurred: {str(e)}"
        print(error_message)
        raise gr.Error(error_message)

# --- Gradio Interface Definition ---
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# Gemini Image Generation Studio")
    gr.Markdown("Use the tabs below to define your image, then click 'Generate Image' to call the API.")

    with gr.Row():
        with gr.Column(scale=1):
            # --- Left Column for Inputs ---
            with gr.Tabs():
                with gr.TabItem("🔑 API & Image"):
                    api_key_input = gr.Textbox(label="Gemini API Key", type="password", info="Your API key is required to generate images.")
                    reference_image_input = gr.Image(label="Reference Image", type="filepath", info="Upload the base image for generation or editing.")

                with gr.TabItem("🎨 Scene & Subject"):
                    scene_input = gr.Textbox(label="Scene", value="cinematic outdoor portrait; professional photography")
                    subject_type_input = gr.Textbox(label="Subject Type", value="adult woman (idol vibe)")
                    age_range_input = gr.Textbox(label="Age Range", value="20s")
                    hair_input = gr.Textbox(label="Hair", value="straight or styled natural open hair with natural shine")
                    makeup_input = gr.Textbox(label="Makeup", value="glossy lips, soft eyeliner, luminous skin")
                    jewellery_input = gr.Textbox(label="Jewellery", value="small hoops, thin chain, subtle bracelets")

                with gr.TabItem("👕 Wardrobe"):
                    top_input = gr.Textbox(label="Top", value="basic tee or camisole")
                    bottom_input = gr.Textbox(label="Bottom", value="denim shorts or mini skirt")
                    footwear_input = gr.Textbox(label="Footwear", value="sneakers or ankle boots")
                    wardrobe_notes_input = gr.Textbox(label="Wardrobe Notes", value="casual modern look, styled for natural setting")

                with gr.TabItem("🧍 Pose & Framing"):
                    pose_angle_input = gr.Dropdown(label="Pose Angle", choices=["three-quarter", "full body"], value="three-quarter")
                    body_pose_input = gr.Textbox(label="Body Pose", value="standing or walking casually, relaxed natural posture")
                    hands_pose_input = gr.Textbox(label="Hands Pose", value="one resting by side or touching hair, the other relaxed")
                    framing_input = gr.Dropdown(label="Framing", choices=["head-to-toe", "waist-up", "cinematic composition"], value="waist-up")

                with gr.TabItem("📷 Camera & Look"):
                    camera_device_input = gr.Textbox(label="Camera Device", value="professional cinema camera / DSLR with prime lens")
                    flash_input = gr.Textbox(label="Flash", value="none; natural golden hour light or soft reflectors")
                    orientation_input = gr.Dropdown(label="Orientation", choices=["vertical", "horizontal"], value="vertical")
                    aspect_ratio_input = gr.Dropdown(label="Aspect Ratio", choices=["16:9", "3:2", "4:3", "1:1"], value="16:9")
                    distance_input = gr.Textbox(label="Distance", value="cinematic portrait distance with shallow depth")
                    focus_input = gr.Textbox(label="Focus", value="sharp on subject; soft bokeh background")
                    texture_input = gr.Textbox(label="Texture", value="smooth high-resolution detail")
                    sharpness_input = gr.Textbox(label="Sharpness", value="very high; crisp cinematic clarity")
                    color_input = gr.Textbox(label="Color", value="warm cinematic grading; golden tones and soft contrast")
                    effects_input = gr.Textbox(label="Effects", value="subtle film grain; natural light flares, depth of field")

                with gr.TabItem("🌳 Background & Style"):
                    background_environment_input = gr.Textbox(label="Background Environment", value="nature setting — forest, park, or meadow with soft light")
                    background_props_input = gr.Textbox(label="Background Props", value="none; focus on subject against natural backdrop")
                    style_genre_input = gr.Textbox(label="Style Genre", value="cinematic portrait photography")
                    authenticity_input = gr.Textbox(label="Authenticity", value="natural, elegant, polished")

                with gr.TabItem("👤 Face & Bans"):
                    use_original_structure_input = gr.Checkbox(label="Use Original Face Structure", value=True)
                    face_description_input = gr.Textbox(label="Face Description", value="maintain the same face shape, features, and proportions as in the provided reference image")
                    gr.Markdown("#### Banned Items")
                    with gr.Row():
                        ban_mirror_input = gr.Checkbox(label="Mirror")
                        ban_phone_input = gr.Checkbox(label="Phone")
                        ban_selfie_input = gr.Checkbox(label="Selfie Look")
                        ban_grainy_input = gr.Checkbox(label="Grainy Noise")
                    with gr.Row():
                        ban_harsh_flash_input = gr.Checkbox(label="Harsh Flash")
                        ban_logos_input = gr.Checkbox(label="Logos")
                        ban_nsfw_input = gr.Checkbox(label="NSFW")
                        ban_cropped_feet_input = gr.Checkbox(label="Cropped Feet")

                with gr.TabItem("⚙️ Output & Variants"):
                    output_count_input = gr.Slider(label="Output Count", minimum=1, maximum=4, step=1, value=1)
                    output_size_input = gr.Textbox(label="Output Size", value="1024x1024")
                    safety_input = gr.Dropdown(label="Safety", choices=["strict", "moderate", "none"], value="strict")
                    variant_name_input = gr.Textbox(label="Variant Name", value="cinematic_nature_fullbody")
                    variant_angle_input = gr.Textbox(label="Variant Angle", value="full-body shot in meadow or forest path, subject centered with depth of field")

        with gr.Column(scale=1):
            # --- Right Column for Outputs ---
            generate_button = gr.Button("Generate Image", variant="primary")
            status_text = gr.Textbox(label="Status", interactive=False)
            image_gallery = gr.Gallery(label="Generated Image(s)", show_label=True, elem_id="gallery", columns=[2], rows=[2], object_fit="contain", height="auto")
            json_output = gr.JSON(label="Generated JSON Input")

    # --- Button Click Action ---
    all_inputs = [
        api_key_input, reference_image_input, scene_input, subject_type_input,
        age_range_input, hair_input, makeup_input, jewellery_input, top_input,
        bottom_input, footwear_input, wardrobe_notes_input, pose_angle_input,
        body_pose_input, hands_pose_input, framing_input, camera_device_input,
        flash_input, orientation_input, aspect_ratio_input, distance_input,
        focus_input, texture_input, sharpness_input, color_input, effects_input,
        background_environment_input, background_props_input, style_genre_input,
        authenticity_input, use_original_structure_input, face_description_input,
        ban_mirror_input, ban_phone_input, ban_selfie_input, ban_grainy_input,
        ban_harsh_flash_input, ban_logos_input, ban_nsfw_input,
        ban_cropped_feet_input, output_count_input, output_size_input,
        safety_input, variant_name_input, variant_angle_input
    ]
    
    generate_button.click(
        fn=generate_image,
        inputs=all_inputs,
        outputs=[image_gallery, json_output, status_text],
    )

if __name__ == "__main__":
    demo.launch(debug=True)