app.py

import gradio as gr
import numpy as np
import os
from PIL import Image
import requests
from io import BytesIO
import io
import base64

# الحصول على التوكن من متغير البيئة
hf_token = os.environ.get("HF_TOKEN_API_DEMO")
if not hf_token:
    raise ValueError("HF_TOKEN_API_DEMO is not set. Please set it as an environment variable.")
auth_headers = {"api_token": hf_token}

# تحويل الصورة إلى Base64
def convert_mask_image_to_base64_string(mask_image):
    buffer = io.BytesIO()
    mask_image.save(buffer, format="PNG")
    image_base64_string = base64.b64encode(buffer.getvalue()).decode('utf-8')
    return f",{image_base64_string}"

# تنزيل الصورة من رابط
def download_image(url):
    response = requests.get(url)
    return Image.open(BytesIO(response.content)).convert("RGB")

# استدعاء API الخاص بالإزالة
def eraser_api_call(image_base64_file, mask_base64_file, mask_type):
    url = "http://engine.prod.bria-api.com/v1/eraser"
    payload = {
        "file": image_base64_file,
        "mask_file": mask_base64_file,
        "mask_type": mask_type,
    }
    response = requests.post(url, json=payload, headers=auth_headers)
    if response.status_code != 200:
        raise ValueError(f"API request failed with status code {response.status_code}: {response.text}")
    response = response.json()
    res_image = download_image(response["result_url"])
    return res_image

# دالة التنبؤ
def predict(dict):
    if 'background' not in dict or 'layers' not in dict:
        raise ValueError("Invalid input format. Missing 'background' or 'layers' keys.")
    init_image = Image.fromarray(dict['background'][:, :, :3], 'RGB')
    mask = Image.fromarray(dict['layers'][0][:, :, 3], 'L')
    image_base64_file = convert_mask_image_to_base64_string(init_image)
    mask_base64_file = convert_mask_image_to_base64_string(mask)
    mask_type = "manual"
    gen_img = eraser_api_call(image_base64_file, mask_base64_file, mask_type)
    return gen_img

# CSS مخصص
css = '''
/* محتوى CSS كما هو */
'''

# واجهة Gradio
image_blocks = gr.Blocks(css=css, elem_id="total-container")
with image_blocks as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown("## BRIA Eraser API")
        gr.HTML('''<p>Demo description here...</p>''')
    with gr.Row():
        with gr.Column():
            image = gr.ImageEditor(
                sources=["upload"],
                layers=False,
                transforms=[],
                brush=gr.Brush(colors=["#000000"], color_mode="fixed")
            )
            with gr.Row(elem_id="prompt-container", equal_height=True):
                btn = gr.Button("Erase!", elem_id="run_button")
        with gr.Column():
            image_out = gr.Image(label="Output", elem_id="output-img")
    btn.click(fn=predict, inputs=image, outputs=image_out, api_name='run')
    gr.HTML("<div class='footer'><p>Footer text here...</p></div>")
image_blocks.queue(max_size=25, api_open=False).launch(show_api=False)