default devs

apocas · Nov 8, 2024 · 11038b3 · 11038b3
1 parent 0ea3f3d
commit 11038b3
Show file tree

Hide file tree

Showing 6 changed files with 13 additions and 0 deletions.
diff --git a/app/image/workers/flux1.py b/app/image/workers/flux1.py
@@ -1,5 +1,6 @@
 import base64
 import io
+import os
 import torch
 from diffusers import FluxPipeline
 import gc
@@ -9,6 +10,8 @@
 
 from app.config import RESTAI_DEFAULT_DEVICE
 
+os.environ["CUDA_VISIBLE_DEVICES"]="0,1,2,3"
+
 def flush():
     gc.collect()
     torch.cuda.empty_cache()

diff --git a/app/image/workers/stablediffusion1_xl_lightning.py b/app/image/workers/stablediffusion1_xl_lightning.py
@@ -1,12 +1,14 @@
 import base64
 import io
+import os
 import torch
 from diffusers import StableDiffusionXLPipeline, UNet2DConditionModel, EulerDiscreteScheduler
 from huggingface_hub import hf_hub_download
 from safetensors.torch import load_file
 
 from app.config import RESTAI_DEFAULT_DEVICE
 
+os.environ["CUDA_VISIBLE_DEVICES"]="0,1,2,3"
 
 def worker(prompt, sharedmem):
     base = "stabilityai/stable-diffusion-xl-base-1.0"

diff --git a/app/image/workers/stablediffusion35_large_turbo.py b/app/image/workers/stablediffusion35_large_turbo.py
@@ -1,12 +1,14 @@
 import base64
 import io
+import os
 from diffusers import BitsAndBytesConfig, SD3Transformer2DModel
 from diffusers import StableDiffusion3Pipeline
 import torch
 from transformers import T5EncoderModel
 
 from app.config import RESTAI_DEFAULT_DEVICE
 
+os.environ["CUDA_VISIBLE_DEVICES"]="0,1,2,3"
 
 def worker(prompt, sharedmem):
     model_id = "stabilityai/stable-diffusion-3.5-large-turbo"

diff --git a/app/image/workers/stablediffusion35_medium.py b/app/image/workers/stablediffusion35_medium.py
@@ -1,12 +1,14 @@
 import base64
 import io
+import os
 from diffusers import BitsAndBytesConfig, SD3Transformer2DModel
 from diffusers import StableDiffusion3Pipeline
 import torch
 from transformers import T5EncoderModel
 
 from app.config import RESTAI_DEFAULT_DEVICE
 
+os.environ["CUDA_VISIBLE_DEVICES"]="0,1,2,3"
 
 def worker(prompt, sharedmem):
     model_id = "stabilityai/stable-diffusion-3.5-medium"

diff --git a/app/image/workers/stablediffusion3_medium.py b/app/image/workers/stablediffusion3_medium.py
@@ -1,10 +1,12 @@
 import base64
 import io
+import os
 from diffusers import StableDiffusion3Pipeline
 import torch
 
 from app.config import RESTAI_DEFAULT_DEVICE
 
+os.environ["CUDA_VISIBLE_DEVICES"]="0,1,2,3"
 
 def worker(prompt, sharedmem):
     base = StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3-medium-diffusers", torch_dtype=torch.float16)

diff --git a/app/image/workers/stablediffusionxl.py b/app/image/workers/stablediffusionxl.py
@@ -1,10 +1,12 @@
 import base64
 import io
+import os
 from diffusers import DiffusionPipeline
 import torch
 
 from app.config import RESTAI_DEFAULT_DEVICE
 
+os.environ["CUDA_VISIBLE_DEVICES"]="0,1,2,3"
 
 def worker(prompt, sharedmem):
     base = DiffusionPipeline.from_pretrained(