import gradio as gr
from gradio_client import Client, file
import tempfile
import os

# Initialize the client
client = Client("radames/Enhance-This-HiDiffusion-SDXL")

# Custom preprocessing function to save uploaded image to a temporary file
def preprocess_image(image):
    # Create a temporary directory
    temp_dir = tempfile.mkdtemp()
    # Save the image to a temporary file
    image_path = os.path.join(temp_dir, "input_image.jpg")
    image.save(image_path)
    return image_path

def enhance_image(image, prompt):
    # Preprocess the image
    image_path = preprocess_image(image)
    # Use the client to predict the result
    (padded_image, image_out), padded_image, anyline_image = client.predict(file(image_path), prompt, api_name="/predict")
    return image_out

# Create the Gradio interface
iface = gr.Interface(
    fn=enhance_image,
    inputs=[
        gr.Image(type="pil", label="Input Image"),
        gr.Textbox(lines=2, placeholder="Enter prompt here", label="Prompt")
    ],
    outputs=gr.Image(type="pil", label="Enhanced Image"),
    title="Image Enhancement with Text Prompt",
    description="Upload an image and provide a text prompt to enhance the image using HiDiffusion SDXL."
)

# Launch the interfacegradio_client.file(
iface.launch()