diff --git a/generate.py b/generate.py index 8e1764e..38ff2d0 100644 --- a/generate.py +++ b/generate.py @@ -1338,3 +1338,94 @@ def reshoot_mustache(character="P1"): extra = "The person is wearing a fake mustache in the style of the mustache from image 1" return process_list(l, reshoot_item, character, extra=extra) + + +def faceswap_item(item, character, keep=False): + if isinstance(character, Document): + chracter_url = public_document_url(character) + else: + chracter_url = character + video_url = public_video_url(item) + model = "half-moon-ai/ai-face-swap/faceswapvideo" + data = { + "source_face_url": chracter_url, + "target_video_url": video_url, + } + prompt_hash = hashlib.sha1(("faceswap:"+chracter_url).encode()).hexdigest() + prefix = "/srv/pandora/static/power/cache/%s_%s" % (item.public_id, prompt_hash) + output_ai = "%s/ai.mp4" % (prefix) + handler = fal_client.submit(model, arguments=data) + request_id = handler.request_id + print(request_id) + result = fal_wait_for(model, request_id) + print(result) + output_url = result["video"]["url"] + ox.net.save_url(output_url, output_ai, overwrite=True) + ai = add_ai_variant(item, output_ai, "ai:0:faceswap") + ai.data["model"] = model + ai.save() + if not keep: + shutil.rmtree(prefix) + return ai + + +def faceswap_item_segments(item, keep=False): + if isinstance(item, str): + item = Item.objects.get(public_id=item) + max_duration = 30 + source = item.files.all()[0].data.path + info = ox.avinfo(source) + duration = info["duration"] + segments = get_item_segments(item, max_duration=max_duration) + prompt_hash = hashlib.sha1("faceswap".encode()).hexdigest() + processed = [] + prefix = "/srv/pandora/static/power/cache/%s_%s" % (item.public_id, prompt_hash) + video_segments = fragment_video(source, prefix, segments) + n = 0 + position = 0 + model = "half-moon-ai/ai-face-swap/faceswapvideo" + for segment in segments: + if isinstance(segment, list): + stype, segment = segment + else: + stype = "n" + character = item.annotations.filter( + layer='prompts', start__gte=position, end__gt=position + ).order_by('start').first().value + character = get_character_document(character) + chracter_url = public_document_url(character) + + output = "%s/%06d.mp4" % (prefix, n) + output_ai = "%s/%06d_ai.mp4" % (prefix, n) + segment_duration = segment - position + if os.path.exists(output): + segment_video_url = public_url(output) + data = { + "source_face_url": chracter_url, + "target_video_url": segment_video_url, + } + handler = fal_client.submit(model, arguments=data) + request_id = handler.request_id + print(request_id) + result = fal_wait_for(model, request_id) + print(result) + output_url = result["video"]["url"] + ox.net.save_url(output_url, output_ai, overwrite=True) + ''' + trimmed = "%s/%06d_ai_trimmed.mp4" % (prefix, n) + frames = int(segment_duration * 24) + trim_video(output_ai, trimmed, frames, stype == "c") + processed.append(trimmed) + ''' + processed.append(output_ai) + n += 1 + position = segment + joined_output = "%s/joined.mp4" % (prefix) + join_segments(processed, joined_output) + ai = add_ai_variant(item, joined_output, "ai:0:faceswap") + ai.data["model"] = model + ai.save() + if not keep: + shutil.rmtree(prefix) + return ai +