Create app.py

app.py

@@ -0,0 +1,116 @@
+import os
+import io
+import base64
+import requests
+import numpy as np
+import gradio as gr
+from PIL import Image
+import onnxruntime
+import cv2
+
+# ——————————————————————————————————————————————————————————————
+#  Configuration
+# ——————————————————————————————————————————————————————————————
+
+HF_TOKEN = os.environ["HF_TOKEN_API_DEMO"]
+AUTH_HEADERS = {"api_token": HF_TOKEN}
+BRIA_API_URL = "http://engine.prod.bria-api.com/v1/gen_fill"
+
+# List your local ONNX upscaler model names (without .ort extension)
+UPSCALE_MODELS = ["modelx2", "modelx4"]
+
+
+# ——————————————————————————————————————————————————————————————
+#  Helper Functions
+# ——————————————————————————————————————————————————————————————
+
+def pil_to_base64(img: Image.Image) -> str:
+    """Convert a PIL image to a base64 string prefixed with a comma."""
+    buf = io.BytesIO()
+    img.save(buf, format="PNG")
+    b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
+    return f",{b64}"
+
+def download_pil_image(url: str) -> Image.Image:
+    r = requests.get(url)
+    return Image.open(io.BytesIO(r.content)).convert("RGB")
+
+def gen_fill(image: Image.Image, mask: Image.Image, prompt: str) -> Image.Image:
+    """Call the BRIA Generative Fill API."""
+    payload = {
+        "file": pil_to_base64(image),
+        "mask_file": pil_to_base64(mask),
+        "prompt": prompt,
+        "steps_num": 12,
+        "sync": True,
+    }
+    res = requests.post(BRIA_API_URL, json=payload, headers=AUTH_HEADERS).json()
+    return download_pil_image(res["urls"][0])
+
+def to_onnx_input(img: np.ndarray) -> np.ndarray:
+    img = img[:, :, :3]  # BGR or RGB first three channels
+    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)  # ensure RGB
+    img = img.astype(np.float32) / 255.0
+    img = np.transpose(img, (2, 0, 1))[None, ...]
+    return img
+
+def from_onnx_output(arr: np.ndarray) -> np.ndarray:
+    arr = np.squeeze(arr, axis=0)
+    arr = np.clip(arr, 0, 1) * 255
+    arr = np.transpose(arr, (1, 2, 0)).astype(np.uint8)
+    return arr
+
+def upscale_image(img: Image.Image, model_name: str) -> Image.Image:
+    """Run ONNX upscaler on a PIL image."""
+    model_path = f"models/{model_name}.ort"
+    sess = onnxruntime.InferenceSession(model_path, sess_options=onnxruntime.SessionOptions())
+    inp = to_onnx_input(np.array(img)[:, :, ::-1])  # PIL is RGB, convert to BGR
+    out = sess.run(None, {sess.get_inputs()[0].name: inp})[0]
+    arr = from_onnx_output(out)
+    # The ONNX model outputs BGR; convert back to RGB
+    rgb = cv2.cvtColor(arr, cv2.COLOR_BGR2RGB)
+    return Image.fromarray(rgb)
+
+# ——————————————————————————————————————————————————————————————
+#  Gradio Interface
+# ——————————————————————————————————————————————————————————————
+
+with gr.Blocks(css="""
+    .gradio-container {max-width: 900px;}
+    #run_button {width:100%; height:48px;}
+    #image_editor img {object-fit: contain; width:100%; height:auto;}
+    #output_col img {object-fit: contain; width:100%; height:auto;}
+""") as demo:
+
+    gr.Markdown("## BRIA Generative Fill + ONNX Upscaler")
+    gr.Markdown("1. Upload your image and draw a mask.  2. Enter a prompt.  3. Choose an upscaler and click **Run**.")
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            editor = gr.ImageEditor(
+                label="Input Image & Mask",
+                tool="editor", brush=gr.Brush(color_mode="binary"),
+                height=400
+            )
+            prompt = gr.Textbox(label="Prompt", placeholder="e.g. “Add a sunset sky”")
+            upscaler = gr.Radio(
+                choices=UPSCALE_MODELS,
+                label="Select Upscaler Model",
+                value=UPSCALE_MODELS[0]
+            )
+            btn = gr.Button("Run", elem_id="run_button")
+
+        with gr.Column(scale=1, elem_id="output_col"):
+            output = gr.Image(label="High-Def Output", height=400)
+
+    def run_pipeline(ed_img, txt, model_name):
+        # ed_img is a RGBA numpy array: [:,:,0:3] = image, [:,:,3] = mask
+        pil_in = Image.fromarray(ed_img[:, :, :3], "RGB")
+        pil_mask = Image.fromarray(ed_img[:, :, 3], "L")
+        filled = gen_fill(pil_in, pil_mask, txt)
+        up_img = upscale_image(filled, model_name)
+        return up_img
+
+    btn.click(fn=run_pipeline, inputs=[editor, prompt, upscaler], outputs=[output])
+
+demo.launch()