Upload app.py

app.py

@@ -0,0 +1,234 @@
+import glob
+import json
+import os
+import subprocess
+import tempfile
+import threading
+import time
+from collections import defaultdict
+from datetime import date
+
+import gradio as gr
+import requests
+
+API_KEY = os.environ.get("DRAMA_API_KEY", "")  # Space secret
+BASE = os.environ.get(
+    "DRAMA_API_BASE",
+    "https://studio.syncmonster.ai/neuralgarage/saas/api/v1/clients",
+)
+HEADERS = {"x-api-key": API_KEY}
+MAX_SECONDS = float(os.environ.get("MAX_SECONDS", 10))
+PER_USER_DAILY = int(os.environ.get("PER_USER_DAILY", 3))
+GLOBAL_DAILY = int(os.environ.get("GLOBAL_DAILY", 200))
+POLL_EVERY = 5      # seconds between status checks
+JOB_TIMEOUT = 900   # give up after 15 minutes
+FONT = "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf"
+
+FRIENDLY_ERRORS = {
+    "ERR_004": "No face found in this video. Try a clip with one clear, front-facing face.",
+    "ERR_006": "The demo is out of credits for now. Please try again later.",
+    "ERR_007": "That file looks corrupted. Try exporting it again.",
+    "ERR_008": "Unsupported file type. Use .mp4 or .mov video and .mp3 or .wav audio.",
+    "ERR_009": "This video is too long.",
+}
+
+_usage = defaultdict(int)
+_lock = threading.Lock()
+
+
+# ---------- helpers ----------
+
+def video_seconds(path):
+    out = subprocess.run(
+        ["ffprobe", "-v", "error", "-show_entries", "format=duration",
+         "-of", "csv=p=0", path],
+        capture_output=True, text=True, check=True,
+    )
+    return float(out.stdout.strip())
+
+
+def extract_audio(video, workdir):
+    wav = os.path.join(workdir, "audio.wav")
+    result = subprocess.run(
+        ["ffmpeg", "-y", "-i", video, "-vn", "-acodec", "pcm_s16le", wav],
+        capture_output=True,
+    )
+    if result.returncode != 0 or not os.path.exists(wav):
+        raise gr.Error("This video has no audio track. Upload an audio file as well.")
+    return wav
+
+
+def take_quota(user):
+    today = date.today().isoformat()
+    with _lock:
+        if _usage[("__all__", today)] >= GLOBAL_DAILY:
+            raise gr.Error("The demo has hit today's limit. Try again tomorrow, "
+                           "or get API access at trydrama.io.")
+        if _usage[(user, today)] >= PER_USER_DAILY:
+            raise gr.Error(f"You've used your {PER_USER_DAILY} free runs for today. "
+                           "Get API access at trydrama.io for more.")
+        _usage[("__all__", today)] += 1
+        _usage[(user, today)] += 1
+
+
+def api_error(body):
+    data = body.get("data") or {}
+    code = data.get("error_code") if isinstance(data, dict) else None
+    return FRIENDLY_ERRORS.get(code, body.get("message") or "Something went wrong.")
+
+
+def build_drama(mood, excited, angry):
+    drama = {}
+    if mood > 0:
+        drama["happy"] = round(mood, 2)
+    elif mood < 0:
+        drama["sad"] = round(-mood, 2)
+    if excited > 0:
+        drama["excited"] = round(excited, 2)
+    if angry > 0:
+        drama["angry"] = round(angry, 2)
+    return drama
+
+
+# ---------- DRAMA API calls ----------
+
+def submit_job(video, audio, drama):
+    with open(video, "rb") as v, open(audio, "rb") as a:
+        r = requests.post(
+            f"{BASE}/generate",
+            headers=HEADERS,
+            data={
+                "skip_find_face": "1",
+                "is_drama": "true",
+                "drama_params": json.dumps(drama),
+            },
+            files={
+                "video": (os.path.basename(video), v),
+                "audio": (os.path.basename(audio), a),
+            },
+            timeout=180,
+        )
+    try:
+        body = r.json()
+    except ValueError:
+        raise gr.Error(f"DRAMA API error (HTTP {r.status_code}).")
+    if body.get("status") != "success":
+        raise gr.Error(api_error(body))
+    return body["data"]["request_id"]
+
+
+def job_status(request_id):
+    r = requests.get(f"{BASE}/get-status", headers=HEADERS,
+                     params={"request_id": request_id}, timeout=30)
+    body = r.json()
+    if body.get("status") != "success":
+        raise gr.Error(api_error(body))
+    return str(body["data"].get("status", "")).lower()
+
+
+def download(request_id, workdir):
+    r = requests.get(f"{BASE}/get-file", headers=HEADERS,
+                     params={"request_id": request_id}, timeout=300, stream=True)
+    if "video" not in r.headers.get("Content-Type", ""):
+        raise gr.Error("The result isn't ready to download yet. Please try again.")
+    path = os.path.join(workdir, "drama_raw.mp4")
+    with open(path, "wb") as f:
+        for chunk in r.iter_content(chunk_size=1 << 20):
+            f.write(chunk)
+    return path
+
+
+def watermark(src, workdir):
+    out = os.path.join(workdir, "drama.mp4")
+    vf = (f"drawtext=fontfile={FONT}:text='trydrama.io':"
+          "x=w-tw-20:y=h-th-20:fontsize=h/24:fontcolor=white@0.85:"
+          "box=1:boxcolor=black@0.35:boxborderw=8")
+    result = subprocess.run(
+        ["ffmpeg", "-y", "-i", src, "-vf", vf, "-c:a", "copy", out],
+        capture_output=True,
+    )
+    return out if result.returncode == 0 else src
+
+
+# ---------- main handler ----------
+
+def run(video, audio, mood, excited, angry, profile: gr.OAuthProfile | None):
+    if not API_KEY:
+        raise gr.Error("The demo isn't switched on yet. Please check back soon.")
+    if profile is None:
+        raise gr.Error("Sign in with Hugging Face (button at the top) to run the demo.")
+    if not video:
+        raise gr.Error("Upload a video first.")
+    if video_seconds(video) > MAX_SECONDS + 0.5:
+        raise gr.Error(f"Please use a clip of {int(MAX_SECONDS)} seconds or less.")
+    drama = build_drama(mood, excited, angry)
+    if not drama:
+        raise gr.Error("Move at least one emotion slider away from 0.")
+
+    take_quota(profile.username)
+    workdir = tempfile.mkdtemp()
+    if not audio:
+        audio = extract_audio(video, workdir)
+
+    yield None, "Uploading your clip…"
+    request_id = submit_job(video, audio, drama)
+
+    started = time.time()
+    while True:
+        status = job_status(request_id)
+        if status == "success":
+            break
+        if status in ("failed", "invalid_request", "internal_error"):
+            raise gr.Error("Processing failed. Try a different clip.")
+        if time.time() - started > JOB_TIMEOUT:
+            raise gr.Error("This is taking too long. Please try again later.")
+        waited = int(time.time() - started)
+        yield None, f"Processing… {waited}s (usually takes a minute or two)"
+        time.sleep(POLL_EVERY)
+
+    yield None, "Downloading the result…"
+    result = watermark(download(request_id, workdir), workdir)
+    yield result, f"Done. Settings used: `{json.dumps(drama)}`"
+
+
+# ---------- UI ----------
+
+with gr.Blocks(title="DRAMA: facial expression editing") as demo:
+    gr.Markdown(
+        "# DRAMA\n"
+        "Change the emotion on a face in any video. Upload a short clip with one "
+        f"clear face (max {int(MAX_SECONDS)} s), set the emotions, and generate.\n\n"
+        "Only upload footage you have the right to use. "
+        "[Model card](https://huggingface.co/trydrama/DRAMA) · "
+        "[API docs](https://docs.drama.syncmonster.ai/) · "
+        "[trydrama.io](https://trydrama.io/?utm_source=huggingface&utm_medium=space)"
+    )
+    gr.LoginButton()
+    with gr.Row():
+        with gr.Column():
+            video_in = gr.Video(label="Video (.mp4 or .mov)", sources=["upload"])
+            audio_in = gr.Audio(
+                label="Audio (optional, defaults to the video's own sound)",
+                type="filepath", sources=["upload"],
+            )
+            mood = gr.Slider(-1, 1, value=0.7, step=0.05,
+                             label="Mood: sad (−1) to happy (+1)")
+            excited = gr.Slider(0, 1, value=0, step=0.05, label="Excited")
+            angry = gr.Slider(0, 1, value=0, step=0.05, label="Angry")
+            go = gr.Button("Generate", variant="primary")
+        with gr.Column():
+            video_out = gr.Video(label="DRAMA result")
+            status_out = gr.Markdown()
+
+    samples = [[p] for p in sorted(glob.glob("examples/*.mp4"))]
+    if samples:
+        gr.Examples(samples, inputs=[video_in], label="Try a sample clip")
+
+    go.click(
+        run,
+        inputs=[video_in, audio_in, mood, excited, angry],
+        outputs=[video_out, status_out],
+        concurrency_limit=2,
+    )
+
+demo.queue(max_size=20).launch()