Update app.py

app.py

@@ -1,7 +1,7 @@
 # ────────────────────────────────────────────────────────────────────────────
 # app.py – Text ➜ 2D (FLUX-mini Kontext) ➜ 3D (Hunyuan3D-2)
 #   • Fits into 16 GB system RAM: 경량 모델 + lazy loading + offload
-#   • Requires: GPU (A10G 24 GB ideal, T4 16 GB OK with fp16)
+#   • Updated: use device_map="balanced" ("auto" not supported by Flux pipelines)
 # ────────────────────────────────────────────────────────────────────────────
 import os
 import tempfile
@@ -27,7 +27,6 @@ DTYPE = torch.float16 if torch.cuda.is_available() else torch.float32
 
 # ─────────────────────── Lazy loaders ───────────────────────
 from diffusers import FluxKontextPipeline, FluxPipeline
-from accelerate import init_empty_weights, load_checkpoint_and_dispatch
 
 # Global caches
 kontext_pipe = None  # type: FluxKontextPipeline | None
@@ -38,16 +37,17 @@ paint_pipe = None
 MINI_KONTEXT_REPO = "black-forest-labs/FLUX.1-Kontext-mini"
 MINI_T2I_REPO = "black-forest-labs/FLUX.1-mini"
 HUNYUAN_REPO = "tencent/Hunyuan3D-2"
+DEVICE_MAP_STRATEGY = "balanced"  # "auto" unsupported for Flux pipelines
 
 
 def load_kontext() -> FluxKontextPipeline:
     global kontext_pipe
     if kontext_pipe is None:
-        print("[+] Loading FLUX.1-Kontext-mini … (low_cpu_mem_usage)")
+        print("[+] Loading FLUX.1-Kontext-mini … (balanced offload)")
         kontext_pipe = FluxKontextPipeline.from_pretrained(
             MINI_KONTEXT_REPO,
             torch_dtype=DTYPE,
-            device_map="auto",
+            device_map=DEVICE_MAP_STRATEGY,
             low_cpu_mem_usage=True,
         )
         kontext_pipe.set_progress_bar_config(disable=True)
@@ -55,14 +55,13 @@ def load_kontext() -> FluxKontextPipeline:
 
 
 def load_text2img() -> FluxPipeline:
-    """Lazy-load light text→image model only when 필요."""
     global _text2img_pipe
     if _text2img_pipe is None:
-        print("[+] Loading FLUX.1-mini (text → image)…")
+        print("[+] Loading FLUX.1-mini (text→image)…")
         _text2img_pipe = FluxPipeline.from_pretrained(
             MINI_T2I_REPO,
             torch_dtype=DTYPE,
-            device_map="auto",
+            device_map=DEVICE_MAP_STRATEGY,
             low_cpu_mem_usage=True,
         )
         _text2img_pipe.set_progress_bar_config(disable=True)
@@ -79,7 +78,7 @@ def load_hunyuan() -> tuple:
         shape_pipe = Hunyuan3DDiTFlowMatchingPipeline.from_pretrained(
             HUNYUAN_REPO,
             torch_dtype=DTYPE,
-            device_map="auto",
+            device_map=DEVICE_MAP_STRATEGY,
             low_cpu_mem_usage=True,
         )
         shape_pipe.set_progress_bar_config(disable=True)
@@ -87,7 +86,7 @@ def load_hunyuan() -> tuple:
         paint_pipe = Hunyuan3DPaintPipeline.from_pretrained(
             HUNYUAN_REPO,
             torch_dtype=DTYPE,
-            device_map="auto",
+            device_map=DEVICE_MAP_STRATEGY,
             low_cpu_mem_usage=True,
         )
         paint_pipe.set_progress_bar_config(disable=True)
@@ -100,7 +99,6 @@ def load_hunyuan() -> tuple:
 def generate_single_2d(prompt: str, image: Image.Image | None, guidance_scale: float) -> Image.Image:
     kontext = load_kontext()
     if image is None:
-        # 텍스트→이미지 : 경량 text2img 파이프라인 사용
         t2i = load_text2img()
         result = t2i(prompt=prompt, guidance_scale=guidance_scale).images[0]
     else:
@@ -116,7 +114,7 @@ def generate_multiview(prompt: str, base_image: Image.Image, guidance_scale: flo
         kontext(image=base_image, prompt=f"{prompt}, right side view", guidance_scale=guidance_scale).images[0],
         kontext(image=base_image, prompt=f"{prompt}, back view", guidance_scale=guidance_scale).images[0],
     ]
-    return views  # [front, left, right, back]
+    return views
 
 
 def build_3d_mesh(prompt: str, images: List[Image.Image]) -> str:
@@ -175,4 +173,3 @@ def build_ui():
 
 if __name__ == "__main__":
     build_ui().queue(max_size=3).launch()
-