import gradio as gr
from PIL import Image
import json

def extract_parameters(image):
    """
    Extract metadata/parameters from a ComfyUI-generated image.
    Returns a formatted string with the extracted info.
    """
    try:
        # Open the uploaded image
        img = Image.open(image)
        # Get metadata from the image's info dictionary
        metadata = img.info
        
        # Check common ComfyUI metadata keys (adjust based on your images)
        if "prompt" in metadata:
            prompt = metadata["prompt"]
            # If it's a JSON string, try to parse it
            try:
                prompt_data = json.loads(prompt)
                # Format JSON nicely
                output = "Extracted Parameters:\n"
                for key, value in prompt_data.items():
                    output += f"{key}: {value}\n"
                return output
            except json.JSONDecodeError:
                # If not JSON, return raw prompt
                return f"Prompt: {prompt}"
        elif metadata:
            # If no "prompt" key, dump all metadata
            output = "Metadata Found:\n"
            for key, value in metadata.items():
                output += f"{key}: {value}\n"
            return output
        else:
            return "No parameters found in image metadata."
    except Exception as e:
        return f"Error processing image: {str(e)}"

# Define the Gradio interface
interface = gr.Interface(
    fn=extract_parameters,
    inputs=gr.Image(type="filepath", label="Upload a ComfyUI Image"),
    outputs=gr.Textbox(label="Extracted Parameters", lines=10),
    title="ComfyUI Image to Parameters",
    description="Upload an image generated by ComfyUI to extract its embedded parameters or prompt.",
    theme="default"
)

# Launch the app
interface.launch()