Update app.py

app.py

@@ -1,91 +1,60 @@
-import insightface
 import os
-import onnxruntime
-import cv2
-import gfpgan
-import tempfile
-import time
+from gradio_client import Client, handle_file
 import gradio as gr
 
-
-class Predictor:
-    def __init__(self):
-        self.setup()
-
-    def setup(self):
-        os.makedirs('models', exist_ok=True)
-        os.chdir('models')
-        if not os.path.exists('GFPGANv1.4.pth'):
-            os.system(
-                'wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth'
-            )
-        if not os.path.exists('inswapper_128.onnx'):
-            os.system(
-                'wget https://huggingface.co/ashleykleynhans/inswapper/resolve/main/inswapper_128.onnx'
-            )
-        os.chdir('..')
-
-        self.face_swapper = insightface.model_zoo.get_model(
-            'models/inswapper_128.onnx',
-            providers=onnxruntime.get_available_providers()
-        )
-        self.face_enhancer = gfpgan.GFPGANer(model_path='models/GFPGANv1.4.pth', upscale=1)
-        self.face_analyser = insightface.app.FaceAnalysis(name='buffalo_l')
-        self.face_analyser.prepare(ctx_id=0, det_size=(640, 640))
-
-    def get_face(self, img_data):
-        analysed = self.face_analyser.get(img_data)
-        try:
-            largest = max(analysed, key=lambda x: (x.bbox[2] - x.bbox[0]) * (x.bbox[3] - x.bbox[1]))
-            return largest
-        except:
-            print("Nuk u gjet fytyra")
-            return None
-
-    def predict(self, input_image, swap_image):
-        try:
-            frame = cv2.imread(input_image.name)
-            face = self.get_face(frame)
-            source_face = self.get_face(cv2.imread(swap_image.name))
-            if not face or not source_face:
-                return None
-            result = self.face_swapper.get(frame, face, source_face, paste_back=True)
-
-            _, _, result = self.face_enhancer.enhance(result, paste_back=True)
-
-            out_path = tempfile.mkdtemp() + f"/{str(int(time.time()))}.jpg"
-            cv2.imwrite(out_path, result)
-            return out_path
-        except Exception as e:
-            print(f"{e}")
-            return None
-
-
-predictor = Predictor()
-
-with gr.Blocks(title="Face Swap & Enhance") as demo:
-    gr.Markdown("## Face Swap & Enhancement Tool")
-    with gr.Row():
-        input_img = gr.Image(type="file", label="Foto tjeter")
-        swap_img = gr.Image(type="file", label="Foto e juaj")
-
-    with gr.Row():
-        output_img = gr.Image(type="filepath", label="Rezultati", visible=False)
-        download_btn = gr.File(label="Shkarko", visible=False)
-
-    submit_btn = gr.Button("Përpunoni")
-
-    def predict_and_update_ui(img1, img2):
-        result = predictor.predict(img1, img2)
-        if result:
-            return result, result, gr.update(visible=True), gr.update(visible=True)
-        else:
-            return None, None, gr.update(visible=False), gr.update(visible=False)
-
-    submit_btn.click(
-        fn=predict_and_update_ui,
-        inputs=[input_img, swap_img],
-        outputs=[output_img, download_btn, output_img, download_btn]
+HF_TOKEN = os.getenv("HF_TOKEN")
+
+client = Client(
+    "franciszzj/Leffa",
+    hf_token=HF_TOKEN,
+)
+
+def virtual_tryon(person_path, garment_path, garment_type):
+    person_file = handle_file(person_path)
+    garment_file = handle_file(garment_path)
+
+    result = client.predict(
+        person_file,
+        garment_file,
+        ref_acceleration=False,
+        step=30,
+        scale=2.5,
+        seed=42,
+        vt_model_type="viton_hd",
+        vt_garment_type=garment_type,
+        vt_repaint=False,
+        api_name="/leffa_predict_vt",
     )
+    return result[0], result[0]  # show image and prepare file for download
 
-demo.launch()
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("")
+    with gr.Row():
+        with gr.Column():
+            gr.Markdown("")
+            src = gr.Image(sources="upload", type="filepath", label="Foto e personit")
+        with gr.Column():
+            gr.Markdown("")
+            ref = gr.Image(sources="upload", type="filepath", label="Foto e veshjes")
+        with gr.Column():
+            gr.Markdown("")
+            garment_type = gr.Radio(
+                choices=[("Siper", "upper_body"),
+                         ("Posht", "lower_body"),
+                         ("Komplet", "dresses")],
+                value="upper_body",
+                label="Lloji i veshjes",
+            )
+        with gr.Column():
+            gr.Markdown("")
+            out = gr.Image(type="filepath", label="Rezultati")
+            download_btn = gr.File(label="Shkarko", interactive=True)
+            btn = gr.Button("VISHU")
+
+    btn.click(virtual_tryon, [src, ref, garment_type], [out, download_btn])
+
+demo.launch(
+    share=True,
+    show_error=True,
+    pwa=True,
+)