Dame bot

bot.py

import asyncio
import tempfile
import discord
from discord.ext import commands
from discord.ext.alternatives import asset_converter, literal_converter

from urllib.parse import urlparse
from typing import Literal
import os


bot = commands.Bot(command_prefix="dame!")

mapping = {
    "dame": ("driver.mp4", "dame.mp3"),
    "rick": ("rick_rendered.mp4", "rick.mp3"),
    "limmy": ("limmy_rendered.mp4", "limmy.mp3"),
    "steel": ("steel_rendered.mp4", "steel.mp3"),
    "potions": ("potions_rendered.mp4", "potions.mp3"),
    "torture": ("torture_rendered.mp4", "torture.mp3")
}

class DeepFake(commands.Cog):
    @commands.max_concurrency(1)
    @commands.command()
    async def df(self, ctx, name: Literal["dame", "rick", "limmy", "steel", "potions", "torture"], image: discord.Asset):
        """Generate a deepfake thing, valid driving videos are: dame rick limmy steel potions torture
        """
        (driver, audio) = mapping[name]

        async with ctx.typing():
            with tempfile.TemporaryDirectory() as d:
                filename = os.path.basename(urlparse(str(image)).path)
                src_name = f"{d}/src_{filename}"
                tmp_name = f"{d}/tmp.mp4"
                dst_name = f"{d}/dst.mp4"

                with open(src_name, "wb+") as src:
                    await image.save(src)

                proc = await asyncio.create_subprocess_exec(
                    "/usr/bin/docker", "run", "--rm", "--gpus", "all", "-v" "/tmp:/tmp", "-v", "/home/ben/dev/first-order-model:/app",
                    "first-order-model:latest", "python3", "demo.py", "--config", "config/vox-adv-256.yaml",
                    "--driving_video", driver,
                    "--source_image", src_name,
                    "--checkpoint", "vox-adv-cpk.pth.tar",
                    "--result_video", tmp_name,
                    "--relative", "--adapt_scale"
                )

                await proc.wait()

                proc = await asyncio.create_subprocess_exec(
                    "/usr/bin/ffmpeg", "-i", tmp_name, "-i", audio, "-c:v", "copy", "-c:a", "aac", dst_name
                )

                await proc.wait()

                with open(dst_name, "rb") as dst:
                    await ctx.send(file=discord.File(dst, filename=f"{name}.mp4"))

    async def cog_command_error(self, ctx, error):
        await ctx.send(error)


bot.add_cog(DeepFake())
bot.run("")

pyproject.toml

[tool.poetry]
name = "first-order-model"
version = "0.1.0"
description = ""
authors = ["Ben Simms <[email protected]>"]

[tool.poetry.dependencies]
python = "^3.8"
"discord.py" = "^1.3.4"
discord-ext-alternatives = { git = "[email protected]:nitros12/discord-ext-alternatives.git", rev = "e5b5990c316ac059f75ecd4138c2ba0ed0249e63" }

[tool.poetry.dev-dependencies]

[build-system]
requires = ["poetry>=0.12"]
build-backend = "poetry.masonry.api"