aiqtech commited on
Commit
ed12c55
·
verified ·
1 Parent(s): c1d518d

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +10 -18
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 by default; GPU if available)
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
- # - Uses fp16 autocast only on CUDA
146
- # - Ensures dtype/device consistency to avoid NaNs
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
  # ---------------------------