Update app.py

app.py

@@ -1,11 +1,11 @@
 # ────────────────────────────────────────────────────────────────────────────
-# app.py – Text ➜ 2D (FLUX-mini Kontext) ➜ 3D (Hunyuan3D-2)
-#   • Fits into 16 GB system RAM: 경량 모델 + lazy loading + offload
-#   • Updated: use device_map="balanced" ("auto" not supported by Flux pipelines)
+# app.py – Text ➜ 2D (FLUX-mini Kontext-dev) ➜ 3D (Hunyuan3D-2)
+#   • Fits into ≈16 GB system RAM: lightweight models + lazy loading + offload
+#   • 2025-07-07: fixed repo names, added HF token + trust_remote_code, cleaned logs
 # ────────────────────────────────────────────────────────────────────────────
 import os
 import tempfile
-from typing import List, Tuple
+from typing import List
 
 import gradio as gr
 import torch
@@ -17,7 +17,7 @@ HF_TOKEN = os.getenv("HF_TOKEN")
 if not HF_TOKEN:
     raise RuntimeError(
         "HF_TOKEN이 설정되지 않았습니다. Space Settings → Secrets에서 "
-        "HF_TOKEN=your_read_token 을 등록한 뒤 재시작하세요."
+        "HF_TOKEN=<your_read_token> 을 등록한 뒤 재시작하세요."
     )
 login(token=HF_TOKEN, add_to_git_credential=False)
 
@@ -29,46 +29,56 @@ DTYPE = torch.float16 if torch.cuda.is_available() else torch.float32
 from diffusers import FluxKontextPipeline, FluxPipeline
 
 # Global caches
-kontext_pipe = None  # type: FluxKontextPipeline | None
-_text2img_pipe = None  # type: FluxPipeline | None
+kontext_pipe: FluxKontextPipeline | None = None
+_text2img_pipe: FluxPipeline | None = None
 shape_pipe = None
 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
+# Repository names (공개 버전)
+MINI_KONTEXT_REPO = "black-forest-labs/FLUX.1-Kontext-dev"   # 이미지 편집/확장용
+MINI_T2I_REPO     = "black-forest-labs/FLUX.1-schnell"       # 텍스트→이미지(4-step distilled)
+HUNYUAN_REPO      = "tencent/Hunyuan3D-2"                    # 3D shape & paint
 
+DEVICE_MAP_STRATEGY = "balanced"  # "auto"(offload) 미지원, so use "balanced"
+
+# ──────────────────────────── Loaders ────────────────────────────
 
 def load_kontext() -> FluxKontextPipeline:
+    """Lazy-load FLUX.1-Kontext-dev (image-to-image editing)."""
     global kontext_pipe
     if kontext_pipe is None:
-        print("[+] Loading FLUX.1-Kontext-mini … (balanced offload)")
+        print("[+] Loading FLUX.1-Kontext-dev … (balanced offload)")
         kontext_pipe = FluxKontextPipeline.from_pretrained(
             MINI_KONTEXT_REPO,
             torch_dtype=DTYPE,
             device_map=DEVICE_MAP_STRATEGY,
             low_cpu_mem_usage=True,
+            token=HF_TOKEN,
+            trust_remote_code=True,
         )
         kontext_pipe.set_progress_bar_config(disable=True)
     return kontext_pipe
 
 
 def load_text2img() -> FluxPipeline:
+    """Lazy-load FLUX.1-schnell (text-to-image)."""
     global _text2img_pipe
     if _text2img_pipe is None:
-        print("[+] Loading FLUX.1-mini (text→image)…")
+        print("[+] Loading FLUX.1-schnell (text→image)…")
         _text2img_pipe = FluxPipeline.from_pretrained(
             MINI_T2I_REPO,
             torch_dtype=DTYPE,
             device_map=DEVICE_MAP_STRATEGY,
             low_cpu_mem_usage=True,
+            token=HF_TOKEN,
+            trust_remote_code=True,
         )
         _text2img_pipe.set_progress_bar_config(disable=True)
     return _text2img_pipe
 
 
-def load_hunyuan() -> tuple:
+def load_hunyuan():
+    """Lazy-load Hunyuan3D-2 shape & texture pipelines."""
     global shape_pipe, paint_pipe
     if shape_pipe is None or paint_pipe is None:
         print("[+] Loading Hunyuan3D-2 (shape & texture)…")
@@ -80,6 +90,8 @@ def load_hunyuan() -> tuple:
             torch_dtype=DTYPE,
             device_map=DEVICE_MAP_STRATEGY,
             low_cpu_mem_usage=True,
+            token=HF_TOKEN,
+            trust_remote_code=True,
         )
         shape_pipe.set_progress_bar_config(disable=True)
 
@@ -88,40 +100,42 @@ def load_hunyuan() -> tuple:
             torch_dtype=DTYPE,
             device_map=DEVICE_MAP_STRATEGY,
             low_cpu_mem_usage=True,
+            token=HF_TOKEN,
+            trust_remote_code=True,
         )
         paint_pipe.set_progress_bar_config(disable=True)
     return shape_pipe, paint_pipe
 
-# ───────────────────────────────────────────────
-# Helper functions
-# ──────────────────��────────────────────────────
+# ───────────────────────────── Helpers ─────────────────────────────
 
 def generate_single_2d(prompt: str, image: Image.Image | None, guidance_scale: float) -> Image.Image:
-    kontext = load_kontext()
+    """Generate a single 2D image (txt2img or img2img)."""
     if image is None:
         t2i = load_text2img()
-        result = t2i(prompt=prompt, guidance_scale=guidance_scale).images[0]
-    else:
-        result = kontext(image=image, prompt=prompt, guidance_scale=guidance_scale).images[0]
-    return result
+        return t2i(prompt=prompt, guidance_scale=guidance_scale).images[0]
+
+    kontext = load_kontext()
+    return kontext(image=image, prompt=prompt, guidance_scale=guidance_scale).images[0]
 
 
 def generate_multiview(prompt: str, base_image: Image.Image, guidance_scale: float) -> List[Image.Image]:
+    """Generate 4-view images for better 3D reconstruction."""
     kontext = load_kontext()
-    views = [
+    return [
         base_image,
         kontext(image=base_image, prompt=f"{prompt}, left side view", guidance_scale=guidance_scale).images[0],
         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
 
 
 def build_3d_mesh(prompt: str, images: List[Image.Image]) -> str:
+    """Create GLB mesh from single or multi-view images."""
     shape, paint = load_hunyuan()
-    single_or_multi = images if len(images) > 1 else images[0]
-    mesh = shape(image=single_or_multi, prompt=prompt)[0]
-    mesh = paint(mesh, image=single_or_multi)
+    source = images if len(images) > 1 else images[0]
+
+    mesh = shape(image=source, prompt=prompt)[0]
+    mesh = paint(mesh, image=source)  # texture painting
 
     tmpdir = tempfile.mkdtemp()
     out_path = os.path.join(tmpdir, "mesh.glb")
@@ -146,7 +160,7 @@ def workflow(prompt: str, input_image: Image.Image | None, multiview: bool, guid
 def build_ui():
     with gr.Blocks(css=CSS, title="Text ➜ 2D ➜ 3D (mini)") as demo:
         gr.Markdown("# 🌀 텍스트 → 2D → 3D 생성기 (경량 버전)")
-        gr.Markdown("Kontext-mini + Hunyuan3D-2. 16 GB RAM에서도 동작합니다.")
+        gr.Markdown("Kontext-dev + Hunyuan3D-2. 16 GB RAM에서도 동작합니다.")
 
         with gr.Row():
             with gr.Column():