Spaces:
Running
on
Zero
Running
on
Zero
Update app.py
Browse files
app.py
CHANGED
@@ -47,8 +47,7 @@ def fixed_pad(self, *args, **kwargs):
|
|
47 |
ChatGLMTokenizer._pad = fixed_pad
|
48 |
|
49 |
# ---------------------------
|
50 |
-
# Load Kolors components
|
51 |
-
# NOTE: dtype is fp16 on CUDA, fp32 on CPU to avoid NaNs on CPU
|
52 |
# ---------------------------
|
53 |
text_encoder = ChatGLMModel.from_pretrained(
|
54 |
f"{ckpt_dir}/text_encoder",
|
@@ -75,7 +74,6 @@ clip_image_encoder = CLIPVisionModelWithProjection.from_pretrained(
|
|
75 |
torch_dtype=DTYPE,
|
76 |
use_safetensors=True
|
77 |
)
|
78 |
-
# Prefer from_pretrained for config parity
|
79 |
clip_image_processor = CLIPImageProcessor.from_pretrained(
|
80 |
"openai/clip-vit-large-patch14-336"
|
81 |
)
|
@@ -95,19 +93,11 @@ pipe = StableDiffusionXLPipeline(
|
|
95 |
print("Models loaded successfully!")
|
96 |
|
97 |
# ---------------------------
|
98 |
-
# InsightFace helper (CPU
|
99 |
# ---------------------------
|
100 |
class FaceInfoGenerator:
|
101 |
def __init__(self, root_dir: str = "./.insightface/"):
|
102 |
-
providers = ["CPUExecutionProvider"]
|
103 |
-
# Try to prefer CUDA provider if available in runtime
|
104 |
-
try:
|
105 |
-
import onnxruntime as ort
|
106 |
-
if "CUDAExecutionProvider" in ort.get_available_providers():
|
107 |
-
providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
|
108 |
-
except Exception:
|
109 |
-
pass
|
110 |
-
|
111 |
self.app = FaceAnalysis(
|
112 |
name="antelopev2",
|
113 |
root=root_dir,
|
@@ -142,10 +132,9 @@ face_info_generator = FaceInfoGenerator()
|
|
142 |
|
143 |
# ---------------------------
|
144 |
# Inference function
|
145 |
-
# -
|
146 |
-
# -
|
147 |
# ---------------------------
|
148 |
-
@spaces.GPU(duration=120)
|
149 |
def infer(
|
150 |
prompt,
|
151 |
image=None,
|
@@ -159,7 +148,7 @@ def infer(
|
|
159 |
gr.Warning("Please upload an image with a face.")
|
160 |
return None, 0
|
161 |
|
162 |
-
# Detect face (InsightFace)
|
163 |
face_info = face_info_generator.get_faceinfo_one_img(image)
|
164 |
if face_info is None:
|
165 |
raise gr.Error("No face detected. Please upload an image with a clear face.")
|
@@ -179,7 +168,6 @@ def infer(
|
|
179 |
pipe.text_encoder = pipe.text_encoder.to(device, dtype=DTYPE)
|
180 |
pipe.unet = pipe.unet.to(device, dtype=DTYPE)
|
181 |
pipe.face_clip_encoder = pipe.face_clip_encoder.to(device, dtype=DTYPE)
|
182 |
-
|
183 |
face_embeds = face_embeds.to(device, dtype=DTYPE)
|
184 |
|
185 |
# Load IP-Adapter weights (FaceID Plus)
|
@@ -235,6 +223,10 @@ def infer(
|
|
235 |
|
236 |
return result, seed
|
237 |
|
|
|
|
|
|
|
|
|
238 |
# ---------------------------
|
239 |
# Gradio UI
|
240 |
# ---------------------------
|
|
|
47 |
ChatGLMTokenizer._pad = fixed_pad
|
48 |
|
49 |
# ---------------------------
|
50 |
+
# Load Kolors components (dtype fp16 on CUDA, fp32 on CPU)
|
|
|
51 |
# ---------------------------
|
52 |
text_encoder = ChatGLMModel.from_pretrained(
|
53 |
f"{ckpt_dir}/text_encoder",
|
|
|
74 |
torch_dtype=DTYPE,
|
75 |
use_safetensors=True
|
76 |
)
|
|
|
77 |
clip_image_processor = CLIPImageProcessor.from_pretrained(
|
78 |
"openai/clip-vit-large-patch14-336"
|
79 |
)
|
|
|
93 |
print("Models loaded successfully!")
|
94 |
|
95 |
# ---------------------------
|
96 |
+
# InsightFace helper (force CPU provider to avoid CUDA init errors)
|
97 |
# ---------------------------
|
98 |
class FaceInfoGenerator:
|
99 |
def __init__(self, root_dir: str = "./.insightface/"):
|
100 |
+
providers = ["CPUExecutionProvider"] # GPU 없는 환경에서 안전
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
101 |
self.app = FaceAnalysis(
|
102 |
name="antelopev2",
|
103 |
root=root_dir,
|
|
|
132 |
|
133 |
# ---------------------------
|
134 |
# Inference function
|
135 |
+
# - No @spaces.GPU decorator (GPU 없을 때 충돌 방지)
|
136 |
+
# - Autocast only on CUDA
|
137 |
# ---------------------------
|
|
|
138 |
def infer(
|
139 |
prompt,
|
140 |
image=None,
|
|
|
148 |
gr.Warning("Please upload an image with a face.")
|
149 |
return None, 0
|
150 |
|
151 |
+
# Detect face (InsightFace on CPU)
|
152 |
face_info = face_info_generator.get_faceinfo_one_img(image)
|
153 |
if face_info is None:
|
154 |
raise gr.Error("No face detected. Please upload an image with a clear face.")
|
|
|
168 |
pipe.text_encoder = pipe.text_encoder.to(device, dtype=DTYPE)
|
169 |
pipe.unet = pipe.unet.to(device, dtype=DTYPE)
|
170 |
pipe.face_clip_encoder = pipe.face_clip_encoder.to(device, dtype=DTYPE)
|
|
|
171 |
face_embeds = face_embeds.to(device, dtype=DTYPE)
|
172 |
|
173 |
# Load IP-Adapter weights (FaceID Plus)
|
|
|
223 |
|
224 |
return result, seed
|
225 |
|
226 |
+
# If CUDA is available, optionally wrap with spaces.GPU for scheduling
|
227 |
+
if torch.cuda.is_available():
|
228 |
+
infer = spaces.GPU(duration=120)(infer)
|
229 |
+
|
230 |
# ---------------------------
|
231 |
# Gradio UI
|
232 |
# ---------------------------
|