Update app.py

app.py

@@ -28,6 +28,14 @@ from pathlib import Path
 from typing import Any, Dict, List, Optional, Tuple, Union
 
 import gradio as gr
+try:
+    import spaces
+except Exception:
+    class _Dummy:
+        def GPU(self, *a, **k):
+            def deco(fn): return fn
+            return deco
+    spaces = _Dummy()
 from PIL import Image, ImageDraw
 
 # ---- HF Spaces GPU decorator (호환용 더미 포함)
@@ -462,6 +470,99 @@ def step1_preprocess(img: Image.Image,
         base = _redraw(base, redraw_strength, redraw_steps, redraw_guidance)
     p = _save_png(base, OUT/"step1"/"input_preprocessed.png")
     return [base], p, "\n".join(logs)
+def step1_cpu(img, keep_rembg, do_weaponless, weapon_terms):
+    """CPU 단계: rembg + (있으면) DINO/SAM/LaMa로 무기 제거. CUDA 사용 금지"""
+    logs = []
+    if img is None:
+        raise gr.Error("이미지를 업로드하세요.")
+    base = img.convert("RGBA")
+
+    # rembg
+    try:
+        base = _remove_bg(base) if keep_rembg else base
+        if keep_rembg:
+            logs.append("rembg 배경 제거")
+    except Exception as e:
+        logs.append(f"rembg 실패: {e}")
+
+    # DINO/SAM/LaMa 무기 제거
+    try:
+        if do_weaponless:
+            base = _weaponless_pipeline(base, weapon_terms, logs)
+    except Exception as e:
+        logs.append(f"무기 제거 실패: {e}")
+
+    out_path = _save_png(base, OUT / "step1" / "input_preprocessed.png")
+    return [(out_path, "preprocessed")], out_path, "\n".join(logs)
+
+
+@spaces.GPU(duration=600)  # ← ZeroGPU 환경: 여기서만 CUDA/모델 로딩 허용
+def step1_gpu_refine(
+    s1_path,
+    enforce_tpose, tpose_strength, tpose_steps, tpose_guidance,
+    do_redraw_flag, redraw_strength, redraw_steps, redraw_guidance
+):
+    """GPU 단계: ControlNet(OpenPose)로 T-포즈 강제 + img2img 리드로우"""
+    logs = []
+    if not s1_path or not Path(s1_path).exists():
+        raise gr.Error("STEP1 이미지가 없습니다. 먼저 STEP1(CPU)을 실행하세요.")
+    img = Image.open(s1_path).convert("RGBA")
+
+    # ---- T-포즈 (ControlNet/OpenPose)
+    if enforce_tpose:
+        try:
+            from diffusers import ControlNetModel, StableDiffusionControlNetImg2ImgPipeline
+            import torch
+            dev = "cuda" if torch.cuda.is_available() else "cpu"
+            cn = ControlNetModel.from_pretrained(
+                "lllyasviel/control_v11p_sd15_openpose",
+                torch_dtype=(torch.float16 if dev == "cuda" else torch.float32)
+            )
+            pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
+                "runwayml/stable-diffusion-v1-5",
+                controlnet=cn,
+                torch_dtype=(torch.float16 if dev == "cuda" else torch.float32)
+            )
+            if dev == "cuda":
+                pipe.to("cuda")
+            pose_canvas = _draw_tpose_openpose_canvas(size=max(img.size))
+            img = pipe(
+                prompt="T-pose, full body, clean anime lines",
+                image=img.convert("RGB"),
+                control_image=pose_canvas,
+                strength=float(tpose_strength),
+                guidance_scale=float(tpose_guidance),
+                num_inference_steps=int(tpose_steps)
+            ).images[0].convert("RGBA")
+            logs.append("ControlNet(OpenPose) T-포즈 적용")
+        except Exception as e:
+            logs.append(f"T-포즈 ControlNet 실패: {e}")
+
+    # ---- img2img 리드로우 (옵션)
+    if do_redraw_flag:
+        try:
+            from diffusers import StableDiffusionImg2ImgPipeline
+            import torch
+            dev = "cuda" if torch.cuda.is_available() else "cpu"
+            pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
+                "runwayml/stable-diffusion-v1-5",
+                torch_dtype=(torch.float16 if dev == "cuda" else torch.float32)
+            )
+            if dev == "cuda":
+                pipe.to("cuda")
+            img = pipe(
+                prompt="clean anime illustration, sharp lines, simple solid background",
+                image=img.convert("RGB"),
+                strength=float(redraw_strength),
+                guidance_scale=float(redraw_guidance),
+                num_inference_steps=int(redraw_steps)
+            ).images[0].convert("RGBA")
+            logs.append("img2img 리드로우 적용")
+        except Exception as e:
+            logs.append(f"img2img 리드로우 실패: {e}")
+
+    out_path = _save_png(img, OUT / "step1" / "input_preprocessed.png")
+    return [(out_path, "refined")], out_path, "\n".join(logs)
 
 # ---------------------------------
 # STEP2: Spaces call (model + texture)
@@ -674,11 +775,18 @@ with gr.Blocks() as demo:
         s1_path = gr.Textbox(label="STEP1 결과 경로", interactive=False)
         s1_log = gr.Textbox(label="STEP1 로그", interactive=False)
 
-        s1_btn.click(step1_preprocess,
-                     inputs=[s1_img, keep_rembg, do_weaponless, weapon_terms, enforce_tpose,
-                             tpose_strength, tpose_steps, tpose_guidance,
-                             do_redraw_flag, redraw_strength, redraw_steps, redraw_guidance],
-                     outputs=[s1_gallery, s1_path, s1_log])
+        # STEP1 실행: CPU → (then) GPU
+        s1_btn.click(
+            step1_cpu,
+            inputs=[s1_img, keep_rembg, do_weaponless, weapon_terms],
+            outputs=[s1_gallery, s1_path, s1_log]
+        ).then(
+            step1_gpu_refine,
+            inputs=[s1_path, enforce_tpose, tpose_strength, tpose_steps, tpose_guidance,
+                    do_redraw_flag, redraw_strength, redraw_steps, redraw_guidance],
+            outputs=[s1_gallery, s1_path, s1_log]
+        )
+
 
     with gr.Tab("STEP 2 — 3D 생성 (모델/텍스처)"):
         do_texture = gr.Checkbox(value=True, label="텍스처 단계 실행(지원 시)")