Docker.Service.AI.Diffusions.STABLE_DIFFUSION.js

const Security = require('Security');

const Docker = {
    Config: require('./../../../Docker.Config.js'),
    Service: {
        Start: require("./../../Docker.Service.Start"),
        AI: {
            Diffusions: {
                STABLE_DIFFUSION: (serviceName, socketPort, gpuAcceleration = false ) => {
                    serviceName = Security.String(serviceName);
                    socketPort = Security.Integer(socketPort);
                    Security.Boolean(gpuAcceleration);

                    const networkName = Docker.Config._networkName

                    const dockerfileContent = [
                        `FROM python:3.10`,
                        `WORKDIR /app`,
                        `RUN apt-get update && apt-get install -y nano`,
                        `RUN pip install torch websockets diffusers transformers accelerate`,
                        `COPY server.py ./`,
                        `RUN python -c "from diffusers import StableDiffusionPipeline; StableDiffusionPipeline.from_pretrained('CompVis/stable-diffusion-v1-4')"`,
                        `CMD ["python", "server.py"]`
                    ].join("\n");

                    let containerArgs = [
                        `-p${socketPort}:1337`,
                        `--network`, networkName
                    ];
                    if (gpuAcceleration) {
                        containerArgs = containerArgs.concat([`--gpus`, `all`])
                    }

                    const contextPath = __dirname+`/../../../../resources/Docker.Service.AI.Diffusions/Docker.Service.AI.Diffusions.STABLE_DIFFUSION`;
                    return Docker.Service.Start(serviceName, dockerfileContent, contextPath, containerArgs);
                }
            },
        }
    }
}

module.exports = Docker.Service.AI.Diffusions.STABLE_DIFFUSION;