app.py

import gradio as gr
import os
import sys
import base64
import uuid
from fastapi import FastAPI

# --- Purana Code (Koi Change Nahi) ---
sys.path.append('roop')
from roop.core import run as roop_run

def face_swap_logic(source_path, target_path):
    output_filename = os.path.basename(target_path)
    os.makedirs("output", exist_ok=True)
    output_path = os.path.join("output", f"api_result_{output_filename}")

    args = [
        "run.py", "--source", source_path, "--target", target_path,
        "--output", output_path, "--execution-provider", "cpu",
    ]
    sys.argv = args
    print(f"Running roop with args: {sys.argv}")

    try:
        roop_run()
    except SystemExit: pass
    except Exception as e:
        print(f"Error during face swap: {e}")
        return None

    if not os.path.exists(output_path):
        return None
        
    return output_path

# --- Gradio UI (Yeh Web Users ke liye hai) ---
with gr.Blocks(theme=gr.themes.Soft()) as ui_app:
    gr.Markdown("# 🎭 Roop Face Swap AI")
    with gr.Tabs():
        with gr.TabItem("🖼️ Swap on Image"):
            with gr.Row():
                source_face_img = gr.Image(label="Source Face", type="filepath")
                target_image = gr.Image(label="Target Image", type="filepath")
            submit_btn_image = gr.Button("Swap Face on Image")
            result_image = gr.Image(label="Result Image")
        with gr.TabItem("🎬 Swap on Video"):
            with gr.Row():
                source_face_vid = gr.Image(label="Source Face", type="filepath")
                target_video = gr.Video(label="Target Video")
            submit_btn_video = gr.Button("Swap Face on Video")
            result_video = gr.Video(label="Result Video")

    def web_ui_swap(source, target):
        return face_swap_logic(source, target)

    submit_btn_image.click(fn=web_ui_swap, inputs=[source_face_img, target_image], outputs=result_image)
    submit_btn_video.click(fn=web_ui_swap, inputs=[source_face_vid, target_video], outputs=result_video)

# --- Naya API Code (Yeh Aapke Android App ke liye hai) ---
app = FastAPI()

def save_base64_to_file(b64_string, extension):
    filename = f"/tmp/{uuid.uuid4().hex}.{extension}"
    with open(filename, "wb") as f:
        f.write(base64.b64decode(b64_string))
    return filename

@app.post("/api/swap-face")
async def api_face_swap(payload: dict):
    source_b64 = payload.get("source_face_b64")
    target_b64 = payload.get("target_media_b64")
    target_ext = payload.get("target_ext", "jpg")

    if not source_b64 or not target_b64:
        return {"error": "Source ya Target file nahi mili."}

    source_file_path = save_base64_to_file(source_b64, "jpg")
    target_file_path = save_base64_to_file(target_b64, target_ext)

    result_path = face_swap_logic(source_file_path, target_file_path)

    if result_path:
        with open(result_path, "rb") as f:
            result_b64 = base64.b64encode(f.read()).decode("utf-8")
        
        os.remove(source_file_path)
        os.remove(target_file_path)
        os.remove(result_path)

        return {"result_b64": result_b64}
    else:
        return {"error": "Face swap fail hua."}

# ===== YAHAN MERI GALTI THI, AB THEEK HAI =====
# OLD: app = gr.mount_app(app, ui_app, path="/")
# NEW: Sahi function ka naam 'mount_gradio_app' hai
app = gr.mount_gradio_app(app, ui_app, path="/")
# ===============================================