import requests
import base64
from PIL import Image
import io

def call_inpaint_api(image_path, mask_path, prompt):
    # Update with your actual space URL
    url = "https://yaseengoldfinchpc-modeltest.hf.space/inpaint"
    
    files = {
        'image': open(image_path, 'rb'),
        'mask': open(mask_path, 'rb')
    }
    
    data = {
        'prompt': prompt
    }
    
    response = requests.post(url, files=files, data=data)
    
    if response.status_code == 200:
        # Decode base64 image
        img_data = base64.b64decode(response.json()['image'])
        img = Image.open(io.BytesIO(img_data))
        img.save('result.png')
        return 'result.png'
    else:
        print(f"Error: {response.text}")
        return None

def main():
    # First test health endpoint
    print("Testing health endpoint...")
    response = requests.get("https://yaseengoldfinchpc-modeltest.hf.space/health")
    print("Health Status:", response.json())
    
    # Test inpainting
    print("\nTesting inpainting...")
    # Replace these with your actual image paths
    image_path = r"C:\Users\M. Y\Downloads\t2.png"  # Replace with your image path
    mask_path = "generated_mask_1.png"   # Replace with your mask path
    prompt = "add some Chair tables and a fountain"
    
    result = call_inpaint_api(image_path, mask_path, prompt)
    if result:
        print(f"Success! Result saved as: {result}")

if __name__ == "__main__":
    main()