Upload app.py with huggingface_hub

app.py

@@ -48,6 +48,7 @@ BUDGETS = {
     "Fast · 512": (512, 512 * 896),
 }
 DEFAULT_BUDGET = "Balanced · 640"
+HARMONIZE_DROP_FRAC = 0.5  # harmonize: drop the mask after this fraction of the denoise steps (midpoint)
 OUTPUT_DIR = os.path.join(tempfile.gettempdir(), "bgswap-outputs")
 os.makedirs(OUTPUT_DIR, exist_ok=True)
 
@@ -329,11 +330,10 @@ def get_duration(request, steps, seed, **_):
     sequence = (int(request["text_token_tags"].shape[0]) + reference_rows(references, num_frames)
                 + target_rows(height, width, num_frames))
     denoise = int(steps) * (STEP_LINEAR * sequence + STEP_QUADRATIC * sequence ** 2) * SAFETY
-    passes = 2 if request.get("harmonize") else 1  # harmonize adds a second (mask-drop) denoise + decode pass
     encode = 5 + reference_rows(references, num_frames) * 1e-3
     share = (height * width * num_frames) / DEFAULT_CANVAS_PIXELS
-    total = (PLACEMENT_ALLOWANCE + encode + SOURCE_ENCODE_PER_DEFAULT_CANVAS * share + passes * denoise
-             + passes * (DECODE_BASE + DECODE_PER_DEFAULT_CANVAS * share) + 10)
+    total = (PLACEMENT_ALLOWANCE + encode + SOURCE_ENCODE_PER_DEFAULT_CANVAS * share + denoise
+             + DECODE_BASE + DECODE_PER_DEFAULT_CANVAS * share + 10)
     duration = max(MIN_GPU_DURATION, min(MAX_GPU_DURATION, int(total)))
     print(f"[bgswap] S={sequence} -> reserving {duration}s ({denoise:.0f}s denoise at {steps} steps)", flush=True)
     return duration
@@ -381,31 +381,21 @@ def _generate(request, steps, seed):
     except Exception as error:
         print(f"[bgswap] turbo switch failed ({type(error).__name__}: {error})", flush=True)
 
-    refs = build_references(request["references"])
-    embeds = request["prompt_embeds"].to("cuda")
-
-    def one_pass(source, mask, source_audio, sr):
-        return PIPE(
-            prompt_embeds=embeds, text_token_tags=request["text_token_tags"], references=refs,
-            source_video=source, source_fps=FPS, mask=mask,
-            source_audio=source_audio, source_audio_sample_rate=sr, audio_mask=None,
-            height=request["height"], width=request["width"], num_frames=request["num_frames"],
-            num_inference_steps=int(steps), generator=torch.Generator("cpu").manual_seed(int(seed)))
-
-    state = one_pass(request["source_video"], request["mask"],
-                     request["source_audio"], request["source_audio_sample_rate"])
-    audio, sr = state.get("audio"), state.get("sampling_rate")
-    videos = state.get("videos")[0]
-
+    kwargs = dict(
+        prompt_embeds=request["prompt_embeds"].to("cuda"), text_token_tags=request["text_token_tags"],
+        references=build_references(request["references"]),
+        source_video=request["source_video"], source_fps=FPS, mask=request["mask"],
+        source_audio=request["source_audio"], source_audio_sample_rate=request["source_audio_sample_rate"],
+        audio_mask=None, height=request["height"], width=request["width"], num_frames=request["num_frames"],
+        num_inference_steps=int(steps), generator=torch.Generator("cpu").manual_seed(int(seed)))
     if request.get("harmonize"):
-        # second pass: drop the mask (uniform value everywhere) and free-refine the whole frame, so the subject's
-        # lighting blends into the new scene. The first pass's audio is kept.
-        src2 = np.stack([np.asarray(f) for f in videos])
-        strength = float(request.get("harmonize_strength", 0.35))
-        mask2 = np.full((src2.shape[0], src2.shape[1], src2.shape[2]), strength, dtype=np.float32)
-        videos = one_pass(src2, mask2, None, None).get("videos")[0]
-
-    return videos, None if audio is None else audio[0].cpu(), sr
+        # single pass: keep the mask until mask_drop_step, then free-refine the whole frame so the subject blends
+        # into the scene (the subject is re-noised to that step's level at the handoff). No second pass.
+        kwargs["mask_drop_step"] = max(1, min(int(steps) - 1, round(int(steps) * HARMONIZE_DROP_FRAC)))
+
+    state = PIPE(**kwargs)
+    audio = state.get("audio")
+    return state.get("videos")[0], None if audio is None else audio[0].cpu(), state.get("sampling_rate")
 
 
 def replace_background(clip, prompt, budget, steps, seed, keep_audio, harmonize=True, remove_only=False, progress=None):