josevalim
commented
1 year ago The docker image is not compiled with cuda but we just released a cuda tag. Can you please try it instead?
Closed xrd closed 1 year ago
josevalim
commented
1 year ago The docker image is not compiled with cuda but we just released a cuda tag. Can you please try it instead?
xrd
commented
1 year ago @josevalim Yes, thank you. I saw the recent PR went onto main, so I assumed it would be in the regular docker image. Thank you, I'll try that and report back.
xrd
commented
1 year ago @josevalim Thank you. It works but then runs out of memory.
I tried the example of changing precision here (https://github.com/elixir-nx/bumblebee/issues/101#issuecomment-1344803404) but am not sure if I converted the code correctly. It is not working yet. I'll keep experimenting. Thanks for your assistance!
repository_id = "CompVis/stable-diffusion-v1-4"
{:ok, tokenizer} = Bumblebee.load_tokenizer({:hf, "openai/clip-vit-large-patch14"})
{:ok, clip} =
Bumblebee.load_model({:hf, repository_id, subdir: "text_encoder"},
log_params_diff: false
)
{:ok, unet} =
Bumblebee.load_model({:hf, repository_id, subdir: "unet"},
params_filename: "diffusion_pytorch_model.bin",
log_params_diff: false
)
##
## Added according to comment in bumblee issue 101)
##
policy = Axon.MixedPrecision.create_policy(compute: :f16)
{:ok, vae} =
Bumblebee.load_model({:hf, repository_id, subdir: "vae"},
architecture: :decoder,
params_filename: "diffusion_pytorch_model.bin",
log_params_diff: false
)
##
## Added according to comment in bumblee issue 101)
##
vae = %{vae | model: Axon.MixedPrecision.apply_policy(clip, policy)}
{:ok, scheduler} = Bumblebee.load_scheduler({:hf, repository_id, subdir: "scheduler"})
{:ok, featurizer} =
Bumblebee.load_featurizer({:hf, repository_id, subdir: "feature_extractor"})
{:ok, safety_checker} =
Bumblebee.load_model({:hf, repository_id, subdir: "safety_checker"},
log_params_diff: false
)
serving =
Bumblebee.Diffusion.StableDiffusion.text_to_image(clip, unet, vae, tokenizer, scheduler,
num_steps: 3,
num_images_per_prompt: 1,
safety_checker: safety_checker,
safety_checker_featurizer: featurizer,
compile: [batch_size: 1, sequence_length: 5],
defn_options: [compiler: EXLA]
)Yeah, we have done zero optimizations, so we hope there is a bunch to gain as we explore it!
I'm trying to use the latest docker image to run a neural network example on GPU/CUDA.
Environment
Erlang/OTP 24 [erts-12.3.2.2] [source] [64-bit] [smp:12:12] [ds:12:12:10] [async-threads:1] [jit] [x86_64-pc-linux-gnu]git rev-parse HEADif running with mix): v0.8.0Current behavior
Start docker with GPUs permitted:
Go to http://localhost:8080.
Go to settings, set
XLA_TARGETtocuda118.Create a new notebook. Add the Mix install code:
It appears to download the CUDA enabled version of XLA for Linux inside the debug output
Once
:okis displayed, then add a smart cell with "Neural Network." Selecttxt2image. Click evaluate.It sees the GPU, and then fails with:
On the console we see this:
Expected behavior
It should let me use the GPU.
NB: I do have a GPU:
Also, I can run other docker images against the GPU: