Update app.py

app.py

@@ -1,4 +1,4 @@
-# app.py (Hugging Face Spaces အတွက် ပြင်ဆင်ပြီး)
+# app.py (Complete Version for Hugging Face Spaces)
 
 import os
 import gradio as gr
@@ -11,16 +11,13 @@ import shutil
 import time
 
 # --- Configuration (Hardcoded for Hugging Face Spaces) ---
-# Command-line arguments များကို ဖယ်ရှားပြီး လိုအပ်သော settings များကို တိုက်ရိုက်သတ်မှတ်ပါသည်
-NUM_FACES = 8  # တစ်ပြိုင်နက်တည်း ပြောင်းနိုင်မည့် မျက်နှာအရေအတွက်
-FORCE_CPU = True # Hugging Face free-tier hardware သည် CPU ဖြစ်သောကြောင့် True ဟု သတ်မှတ်ပါသည်
+NUM_FACES = 8
+FORCE_CPU = True
 
 # --- Initialize Refacer ---
-# Refacer ကို သတ်မှတ်ထားသော settings များဖြင့် စတင်ပါသည်
 refacer = Refacer(force_cpu=FORCE_CPU)
 
-# --- Functions (မူရင်းအတိုင်းနီးပါး၊ မလိုအပ်သော temp folder management ကို ဖယ်ရှားထား) ---
-
+# --- Core Functions ---
 def run_image(*vars):
     image_path = vars[0]
     origins = vars[1:(NUM_FACES + 1)]
@@ -28,22 +25,12 @@ def run_image(*vars):
     thresholds = vars[(NUM_FACES * 2) + 1:-2]
     face_mode = vars[-2]
     partial_reface_ratio = vars[-1]
-
     disable_similarity = (face_mode in ["Single Face", "Multiple Faces"])
     multiple_faces_mode = (face_mode == "Multiple Faces")
-
-    faces = []
-    for k in range(NUM_FACES):
-        if destinations[k] is not None:
-            faces.append({
-                'origin': origins[k] if not multiple_faces_mode else None,
-                'destination': destinations[k],
-                'threshold': thresholds[k] if not multiple_faces_mode else 0.0
-            })
-
+    faces = [{'origin': origins[k] if not multiple_faces_mode else None, 'destination': destinations[k], 'threshold': thresholds[k] if not multiple_faces_mode else 0.0} for k in range(NUM_FACES) if destinations[k] is not None]
     return refacer.reface_image(image_path, faces, disable_similarity=disable_similarity, multiple_faces_mode=multiple_faces_mode, partial_reface_ratio=partial_reface_ratio)
 
-def run(*vars):
+def run_video(*vars):
     video_path = vars[0]
     origins = vars[1:(NUM_FACES + 1)]
     destinations = vars[(NUM_FACES + 1):(NUM_FACES * 2) + 1]
@@ -51,37 +38,22 @@ def run(*vars):
     preview = vars[-3]
     face_mode = vars[-2]
     partial_reface_ratio = vars[-1]
-
     disable_similarity = (face_mode in ["Single Face", "Multiple Faces"])
     multiple_faces_mode = (face_mode == "Multiple Faces")
-
-    faces = []
-    for k in range(NUM_FACES):
-        if destinations[k] is not None:
-            faces.append({
-                'origin': origins[k] if not multiple_faces_mode else None,
-                'destination': destinations[k],
-                'threshold': thresholds[k] if not multiple_faces_mode else 0.0
-            })
-
+    faces = [{'origin': origins[k] if not multiple_faces_mode else None, 'destination': destinations[k], 'threshold': thresholds[k] if not multiple_faces_mode else 0.0} for k in range(NUM_FACES) if destinations[k] is not None]
     mp4_path, gif_path = refacer.reface(video_path, faces, preview=preview, disable_similarity=disable_similarity, multiple_faces_mode=multiple_faces_mode, partial_reface_ratio=partial_reface_ratio)
-    return mp4_path, gif_path if gif_path else None
+    return mp4_path, gif_path
 
 def load_first_frame(filepath):
-    if filepath is None:
-        return None
+    if filepath is None: return None
     with imageio.get_reader(filepath) as reader:
         return reader.get_data(0)
 
 def extract_faces_auto(filepath, max_faces=5, isvideo=False):
-    if filepath is None:
-        return [None] * max_faces
-
-    # Video ဖိုင်ကြီးလျှင် performance မကျစေရန် auto-extract ကို ကျော်ပါသည်
-    if isvideo and os.path.getsize(filepath) > 5 * 1024 * 1024:
-        print("Video is too large for auto-extract, skipping face extraction.")
+    if filepath is None: return [None] * max_faces
+    if isvideo and os.path.getsize(filepath) > 10 * 1024 * 1024:
+        print("Video too large for auto-extract.")
         return [None] * max_faces
-
     try:
         frame = load_first_frame(filepath)
         faces = refacer.extract_faces_from_image(frame, max_faces=max_faces)
@@ -92,96 +64,81 @@ def extract_faces_auto(filepath, max_faces=5, isvideo=False):
 
 def toggle_tabs_and_faces(mode):
     if mode == "Single Face":
-        tab_updates = [gr.update(visible=(i == 0)) for i in range(NUM_FACES)]
-        origin_updates = [gr.update(visible=False)] * NUM_FACES
+        return [gr.update(visible=(i == 0)) for i in range(NUM_FACES)] + [gr.update(visible=False)] * NUM_FACES
     elif mode == "Multiple Faces":
-        tab_updates = [gr.update(visible=True)] * NUM_FACES
-        origin_updates = [gr.update(visible=False)] * NUM_FACES
+        return [gr.update(visible=True)] * NUM_FACES + [gr.update(visible=False)] * NUM_FACES
     else: # Faces By Match
-        tab_updates = [gr.update(visible=True)] * NUM_FACES
-        origin_updates = [gr.update(visible=True)] * NUM_FACES
-    return tab_updates + origin_updates
+        return [gr.update(visible=True)] * NUM_FACES + [gr.update(visible=True)] * NUM_FACES
 
 def handle_tif_preview(filepath):
-    if filepath is None:
-        return None
-    
-    # tempfile ကိုအသုံးပြုခြင်းဖြင့် Gradio က temp files များကို သူ့အလိုလို စီမံခန့်ခွဲစေပါသည်
-    with Image.open(filepath) as img:
-        img_rgb = img.convert('RGB')
-        with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as temp_file:
-            img_rgb.save(temp_file.name)
-            return temp_file.name
+    if filepath is None: return None
+    with Image.open(filepath) as img, tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as temp_file:
+        img.convert('RGB').save(temp_file.name)
+        return temp_file.name
 
 # --- UI ---
 theme = gr.themes.Base(primary_hue="blue", secondary_hue="cyan")
 
 with gr.Blocks(theme=theme, title="NeoRefacer - AI Refacer") as demo:
-    # Icon ကို base64 encode လုပ်ပြီး HTML ထဲမှာ တိုက်ရိုက်ထည့်သွင်းခြင်း
     with open("icon.png", "rb") as f:
         icon_data = base64.b64encode(f.read()).decode()
-    icon_html = f'<img src="data:image/png;base64,{icon_data}" style="width:40px;height:40px;margin-right:10px;">'
-
-    with gr.Row():
-        gr.Markdown(f"""
-        <div style="display: flex; align-items: center;">
-        {icon_html}
-        <span style="font-size: 2em; font-weight: bold; color:#2563eb;">NeoRefacer</span>
-        </div>
-        """)
-
-    # --- IMAGE MODE ---
-    with gr.Tab("Image Mode"):
-        with gr.Row():
-            image_input = gr.Image(label="Original image", type="filepath")
-            image_output = gr.Image(label="Refaced image", interactive=False, type="filepath")
-
-        with gr.Row():
-            face_mode_image = gr.Radio(["Single Face", "Multiple Faces", "Faces By Match"], value="Single Face", label="Replacement Mode")
-            partial_reface_ratio_image = gr.Slider(label="Reface Ratio (0 = Full Face, 0.5 = Half Face)", minimum=0.0, maximum=0.5, value=0.0, step=0.1)
-            image_btn = gr.Button("Reface Image", variant="primary")
-
-        origin_image, destination_image, thresholds_image, face_tabs_image = [], [], [], []
-
-        for i in range(NUM_FACES):
-            with gr.Tab(f"Face #{i+1}", visible=(i==0)) as tab:
-                with gr.Row():
-                    origin = gr.Image(label="Face to replace", visible=True)
-                    destination = gr.Image(label="Destination face")
-                threshold = gr.Slider(label="Threshold", minimum=0.0, maximum=1.0, value=0.2)
-            origin_image.append(origin)
-            destination_image.append(destination)
-            thresholds_image.append(threshold)
-            face_tabs_image.append(tab)
-
-        face_mode_image.change(fn=toggle_tabs_and_faces, inputs=[face_mode_image], outputs=face_tabs_image + origin_image)
-        demo.load(fn=lambda: toggle_tabs_and_faces("Single Face"), inputs=None, outputs=face_tabs_image + origin_image)
-
-        image_btn.click(fn=run_image, inputs=[image_input] + origin_image + destination_image + thresholds_image + [face_mode_image, partial_reface_ratio_image], outputs=[image_output])
-        image_input.change(fn=lambda filepath: extract_faces_auto(filepath, max_faces=NUM_FACES), inputs=image_input, outputs=destination_image) # Auto-fill destination faces
-
-    # --- GIF MODE ---
-    with gr.Tab("GIF Mode"):
-        # (UI code is repetitive, so I'll omit it for brevity, assuming it's similar to the original)
-        # The logic is identical to Image Mode, just with different input/output components.
-        # Ensure all instances of `num_faces` are replaced with `NUM_FACES`.
-        # ... (GIF UI Code Here) ...
-        pass
-        
-    # --- TIF/TIFF MODE ---
-    with gr.Tab("TIFF Mode"):
-        # (UI code is repetitive, so I'll omit it for brevity)
-        # Ensure all instances of `num_faces` are replaced with `NUM_FACES`.
-        # ... (TIFF UI Code Here) ...
-        pass
-
-    # --- VIDEO MODE ---
-    with gr.Tab("Video Mode"):
-        # (UI code is repetitive, so I'll omit it for brevity)
-        # Ensure all instances of `num_faces` are replaced with `NUM_FACES`.
-        # ... (Video UI Code Here) ...
-        pass
+    gr.Markdown(f'<div style="display: flex; align-items: center;"><img src="data:image/png;base64,{icon_data}" style="width:40px;height:40px;margin-right:10px;"><span style="font-size: 2em; font-weight: bold; color:#2563eb;">NeoRefacer</span></div>')
+
+    for mode in ["Image", "GIF", "TIFF", "Video"]:
+        with gr.Tab(f"{mode} Mode"):
+            is_video_or_gif = mode in ["Video", "GIF"]
+            
+            with gr.Row():
+                if is_video_or_gif:
+                    input_component = gr.Video(label=f"Original {mode}")
+                    if mode == "GIF":
+                        output_main = gr.Video(label="Refaced GIF (MP4)", interactive=False, format="mp4")
+                        output_secondary = gr.Image(label="Refaced GIF (GIF)", type="filepath")
+                    else: # Video
+                        output_main = gr.Video(label="Refaced Video", interactive=False, format="mp4")
+                        output_secondary = gr.File(visible=False) # Dummy component
+                elif mode == "TIFF":
+                    input_component = gr.File(label="Original TIF", file_types=[".tif", ".tiff"])
+                    output_main = gr.Image(label="Refaced TIF Preview", type="filepath")
+                    output_secondary = gr.File(label="Refaced TIF (Download)", interactive=False)
+                else: # Image
+                    input_component = gr.Image(label="Original image", type="filepath")
+                    output_main = gr.Image(label="Refaced image", interactive=False, type="filepath")
+            
+            with gr.Row():
+                face_mode_radio = gr.Radio(["Single Face", "Multiple Faces", "Faces By Match"], value="Single Face", label="Replacement Mode")
+                partial_reface_slider = gr.Slider(label="Reface Ratio (0=Full, 0.5=Half)", minimum=0.0, maximum=0.5, value=0.0, step=0.1)
+                reface_btn = gr.Button(f"Reface {mode}", variant="primary")
+                if is_video_or_gif:
+                    preview_checkbox = gr.Checkbox(label="Preview (fast)", value=False)
+            
+            origins, destinations, thresholds, face_tabs = [], [], [], []
+            for i in range(NUM_FACES):
+                with gr.Tab(f"Face #{i+1}", visible=(i==0)) as tab:
+                    with gr.Row():
+                        origin_img = gr.Image(label="Face to replace (Match)", visible=False)
+                        dest_img = gr.Image(label="Destination face (Target)")
+                    thresh_slider = gr.Slider(label="Threshold (for Match mode)", minimum=0.0, maximum=1.0, value=0.2)
+                origins.append(origin_img); destinations.append(dest_img); thresholds.append(thresh_slider); face_tabs.append(tab)
+
+            # Event Handlers
+            face_mode_radio.change(toggle_tabs_and_faces, inputs=face_mode_radio, outputs=face_tabs + origins)
+            
+            if mode == "TIFF":
+                tif_preview = gr.Image(label="TIF Preview", type="filepath") # Specific for TIFF
+                input_component.change(handle_tif_preview, inputs=input_component, outputs=tif_preview)
+                input_component.change(lambda fp: extract_faces_auto(fp, max_faces=NUM_FACES), inputs=input_component, outputs=destinations)
+                reface_btn.click(lambda fp, *args: (handle_tif_preview(run_image(fp, *args)), run_image(fp, *args)), inputs=[input_component] + origins + destinations + thresholds + [face_mode_radio, partial_reface_slider], outputs=[output_main, output_secondary])
+            elif is_video_or_gif:
+                input_component.change(lambda fp: extract_faces_auto(fp, max_faces=NUM_FACES, isvideo=True), inputs=input_component, outputs=destinations)
+                reface_btn.click(run_video, inputs=[input_component] + origins + destinations + thresholds + [preview_checkbox, face_mode_radio, partial_reface_slider], outputs=[output_main, output_secondary])
+            else: # Image
+                input_component.change(lambda fp: extract_faces_auto(fp, max_faces=NUM_FACES), inputs=input_component, outputs=destinations)
+                reface_btn.click(run_image, inputs=[input_component] + origins + destinations + thresholds + [face_mode_radio, partial_reface_slider], outputs=[output_main])
+    
+    # Load initial state for all tabs
+    demo.load(lambda: toggle_tabs_and_faces("Single Face"), outputs=face_tabs + origins)
 
 # --- Launch app ---
-# Hugging Face Spaces အတွက် `launch()` ကို argument မထည့်ဘဲ ခေါ်ပါ
-demo.queue().launch()
+demo.queue().launch()
+