Skip to content

Part 2 - Advanced OctoShop Pipeline (ControlNet, Face Preservation)

In this lab we'll augment the OctoShop Pipeline with the use of ControlNet for pose/scene preservation and AI-based face swap for face preservation.

Lab 2 is a bit more advanced than Lab 1 because you'll be expected to make some modifications to source code. This Lab is a bit more open ended than Lab 1.

Background on ControlNet

ControlNet is a neural network structure to control diffusion models by adding extra conditions.

In other words, a ControlNet lets us produce a SDXL generated image that is consistent with a control input image. The point of us building ControlNet into our SDXL model container is to give us an advanced OctoShop workflow that produces output images that are consistent with the user-provided input images.

Consistency is an important concept in Generative AI image or video pipelines. For instance in video, it's important to keep each frame consistent with the precedent frame to deliver a smooth, non jittery video stream.

Section 1: Extend the SDXL model container to support ControlNet

In this section we'll edit the ~/dockercon23-octoai/lab2/sdxl_canny/model.py file to support the Canny-based ControlNet. We've built a scaffold for you based on Lab 1 SDXL model container.

Go ahead and modify the source code to extend the functionality of the model container to support ControlNet-Canny.

Fill in all of the #TODOs left in the model.py file.

Your resources to successfully complete this task are the following:

  • Use the following link from Hugging Face to understand how ControlNet-Canny is being used in SDXL pipeline.
    • Note in the usage code example how the controlnet, vae, pipe are all being initialized.
    • Note also how Canny (edge filter from OpenCV) is being used to pre-process the ControlNet control image.
    • Finally note how the ControlNet-Canny-SDXL pipeline is being called, and what arguments are being passed in.
  • When invoking the StableDiffusionXLControlNetPipeline, note the arguments that you can pass in based on the HuggingFace Diffusers technical documentation, particularly the ControlNet-specific parameters such as:
    • control_guidance_start
    • control_guidance_end
    • controlnet_conditioning_scale

When you are done modifying the code of model.py, you'll need to test it. We'll explain in the next section how you would go about doing this.

Section 2: Building and running the SDXL+ControlNet container

Make the sdxl-canny directory your working directory.

cd ~/dockercon23-octoai/lab2/sdxl_canny

Build the container with the following command

docker build -t sdxl-canny --progress plain .

Check that the sdxl-canny image was built successfully with the following command

docker image ls

You'll see something like this 

REPOSITORY                                    TAG                                    IMAGE ID       CREATED         SIZE
sdxl-canny                                    latest                                 ec2204886c59   5 minutes ago   31.2GB

Next, let's run the container locally

docker run -d --rm -p 8080:8080 \
        --gpus all \
        --name sdxl-canny sdxl-canny

Now that this container is running locally, double check that it is running

docker container ls

You'll see something like this 

CONTAINER ID   IMAGE        COMMAND                  CREATED         STATUS         PORTS                                       NAMES
7a6598aab4e6   sdxl-canny   "/opt/nvidia/nvidia_…"   2 minutes ago   Up 2 minutes   0.0.0.0:8080->8080/tcp, :::8080->8080/tcp   sdxl-canny

You can dump the logs with the following command while the container is running

docker logs sdxl-canny

You'll see the following log run the model initialization code 

Initializing SDXL Canny pipeline...
Loading pipeline components...: 100%|██████████| 7/7 [02:03<00:00, 17.58s/it]
SDXL Canny pipeline has been initialized!

That means the model container is ready to serve inference requests! It'll be listening at the following URL for incoming inference requests: http://localhost:8080/predict.

Section 3: Test, deploy, and integrate the model container into an OctoShop pipeline

Once the model container is running locally, it's time to fire up the lab2 iPython Notebook to test it out.

Just like before, launch the iPython Notebook to test out the SDXL model container.

We'll be launching the jupyter notebook from the AWS instance, since we need to be able to interface with the docker container we just launched on that instance.

cd ~/dockercon23-octoai/lab2
jupyter notebook --no-browser --port=8888

Your AWS dev instance will launch the notebook server and produce the following log. Note the URLs that is provided in the log (your token will be unique and from the one below different, so we put in a place holder): 

    Or copy and paste one of these URLs:
        http://localhost:8888/tree?token=735c08fd361cc0fae396aba2a68667b93657d6798a2e2a76

Then on your laptop browser paste the URL with your own token (provided by Jupyter Notebook server). This assumes you've set up port forwarding between the AWS instance and your laptop as indicated in the "SSH into your instance" of the pre-requisites.

Note: your token will certainly be very different from the one below so make sure to use your own token!

http://localhost:8888/tree?token=735c08fd361cc0fae396aba2a68667b93657d6798a2e2a76

The rest of the instructions for section 3 will be embedded in the lab2 iPython notebook.

When you're done, you can kill the Jupyter notebook server by hitting Ctrl+C from the remote session that had launched the jupyter notebook session.

Section 4: Stop your docker container

When you're done running local tests, you can go ahead and stop the sdxl-canny container that you've been running with the following command

docker stop sdxl-canny

Section 6: Build and run your discord bot

  • Go to discordbot folder
cd ~/dockercon23-octoai/lab2/discordbot
    • Personalize your discord bot by overwriting the following lines in workflow.py, under the octoshop_workflow() function definition.
user_prompt = "set in outer space"

user_style = {
    "name": "sai-digital art",
    "prompt": "concept art {prompt} . digital artwork, illustrative, painterly, matte painting, highly detailed",
    "negative_prompt": "photo, photorealistic, realism, ugly",
}
  • Build container
docker build -t bot .
  • Run container with three variables below. YOUR_COMMAND_NAME is the name of the bot command you will use in the discord server.
    • Name it to whatever you want! As long as it's unique (i.e. doesn't alias with the other contestants). For instance for this bot that puts photos into space, I may want to call my command spaceify. I'll therefore be able to invoke the bot on Discord using the /spacify command.
    • This can be edited right before running the container and it doesn't require any modification in your bot account.
    • Once you've settled on a bot name, please update the workshop attendee spreadsheet under the Discord Bot Cmd column.
docker run --rm -it \
-e DISCORD_TOKEN=YOUR_DISCORD_TOKEN \
-e DISCORD_COMMAND=YOUR_COMMAND_NAME \
-e SDXL_ENDPOINT_URL=YOUR_SDXL_OCTOAI_ENDPOINT_URL \
bot
  • Now that the container is running on your AWS dev instance, test out your bot! Note that your bot will be functional as long as the bot container is running on the AWS instance. This means your bot will be functional for the workshop and likely not beyond the day of the workshop.

When you're done running your bot, you can kill the bot container with Ctrl+C. This will cause the bot to stop functioning.

Section 7: Pause your OctoAI SDXL endpoint

When you don't need to use your OctoAI SDXL endpoint, you can go to the Info view of the SDXL endpoint and simply hit Pause endpoint button to pause the endpoint.

While the endpoint is paused you won't be billed for usage.