Update app.py
Browse files
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 = "/
|
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 |
-
"/
|
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")
|