stzhao commited on
Commit
2efad28
·
verified ·
1 Parent(s): 6c779be

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +4 -4
app.py CHANGED
@@ -4,12 +4,12 @@ import torch
4
  from diffusers import Lumina2Pipeline
5
  from transformers import AutoModelForCausalLM, AutoTokenizer
6
 
7
- # Set up environment
8
- os.environ['CUDA_VISIBLE_DEVICES'] = "0"
9
 
10
  # Load models
11
  def load_models():
12
- model_name = "/mnt/petrelfs/zhaoshitian/models/LeX-Enhancer-full"
13
 
14
  model = AutoModelForCausalLM.from_pretrained(
15
  model_name,
@@ -19,7 +19,7 @@ def load_models():
19
  tokenizer = AutoTokenizer.from_pretrained(model_name)
20
 
21
  pipe = Lumina2Pipeline.from_pretrained(
22
- "/mnt/hwfile/alpha_vl/qilongwu/checkpoints/LeX-Lumina",
23
  torch_dtype=torch.bfloat16
24
  )
25
  pipe.to("cuda")
 
4
  from diffusers import Lumina2Pipeline
5
  from transformers import AutoModelForCausalLM, AutoTokenizer
6
 
7
+ # # Set up environment
8
+ # os.environ['CUDA_VISIBLE_DEVICES'] = "0"
9
 
10
  # Load models
11
  def load_models():
12
+ model_name = "X-ART/LeX-Enhancer-full"
13
 
14
  model = AutoModelForCausalLM.from_pretrained(
15
  model_name,
 
19
  tokenizer = AutoTokenizer.from_pretrained(model_name)
20
 
21
  pipe = Lumina2Pipeline.from_pretrained(
22
+ "X-ART/LeX-Lumina",
23
  torch_dtype=torch.bfloat16
24
  )
25
  pipe.to("cuda")