import gradio as gr
from huggingface_hub import InferenceClient
from PIL import Image
import io

# Initialize the Inference Client
client = InferenceClient()

# Function to handle the image generation
def generate_image(image, prompt):
    # Open the uploaded image
    image = Image.open(image)

    # Generate image using the InferenceClient
    generated_image = client.image_to_image(image, prompt=prompt)

    # Save generated image to a BytesIO object
    img_byte_arr = io.BytesIO()
    generated_image.save(img_byte_arr, format='PNG')
    img_byte_arr.seek(0)  # Move the cursor to the beginning

    # Return the generated image
    return generated_image

# Create Gradio interface
iface = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.inputs.Image(type="file", label="Input Image"),  # File input for uploaded image
        gr.inputs.Textbox(label="Prompt")  # Textbox for the prompt
    ],
    outputs=gr.outputs.Image(type="pil", label="Generated Image"),  # Output as PIL image
    title="Image Generation with Hugging Face",
    description="Upload an image and provide a prompt to generate a new image."
)

# Launch the Gradio interface
iface.launch()