danhtran2mind commited on
Commit
4d7446e
·
verified ·
1 Parent(s): be63645

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +6 -1
app.py CHANGED
@@ -1,7 +1,7 @@
1
  import dataclasses
2
  import json
3
  from pathlib import Path
4
-
5
  import gradio as gr
6
  import torch
7
  from PIL import Image
@@ -54,6 +54,11 @@ def create_demo(
54
  device = torch.device(device)
55
  dtype = torch.float16 if device.type == "cuda" else torch.float32
56
 
 
 
 
 
 
57
  # Load models with consistent dtype
58
  vae = AutoencoderKL.from_pretrained(model_name, subfolder="vae", torch_dtype=dtype).to(device)
59
  tokenizer = CLIPTokenizer.from_pretrained(model_name, subfolder="tokenizer")
 
1
  import dataclasses
2
  import json
3
  from pathlib import Path
4
+ import os
5
  import gradio as gr
6
  import torch
7
  from PIL import Image
 
54
  device = torch.device(device)
55
  dtype = torch.float16 if device.type == "cuda" else torch.float32
56
 
57
+ # Set Number CPU core is maximum
58
+ num_cpu_cores = os.cpu_count() # Adjust based on your system
59
+ os.environ["OMP_NUM_THREADS"] = str(num_cpu_cores)
60
+ torch.set_num_threads(num_cpu_cores)
61
+
62
  # Load models with consistent dtype
63
  vae = AutoencoderKL.from_pretrained(model_name, subfolder="vae", torch_dtype=dtype).to(device)
64
  tokenizer = CLIPTokenizer.from_pretrained(model_name, subfolder="tokenizer")