Update app.py

app.py

@@ -1,8 +1,9 @@
 import spaces
 import time
 import os
+import tempfile
 
-# ONNX Runtime에서 CUDA provider 우선 사용 (insightface → ORT 세션에 반영)
+# ONNX Runtime CUDA provider 시도 (효과 없더라도 무해)
 os.environ.setdefault("INSIGHTFACE_ONNX_PROVIDERS", "CUDAExecutionProvider,CPUExecutionProvider")
 os.environ.setdefault("ORT_LOG_severity_level", "3")  # ORT 로그 최소화
 
@@ -49,6 +50,14 @@ class FluxGenerator:
 flux_generator = FluxGenerator()
 
 
+def _save_pil(img: Image.Image, prefix: str = "out") -> str:
+    os.makedirs("/tmp", exist_ok=True)
+    ts = int(time.time() * 1000)
+    path = f"/tmp/{prefix}_{ts}.png"
+    img.save(path, format="PNG")
+    return path
+
+
 @spaces.GPU
 @torch.inference_mode()
 def generate_image(
@@ -97,8 +106,6 @@ def generate_image(
         id_embeddings = None
         uncond_id_embeddings = None
 
-    print(id_embeddings)
-
     # prepare input
     x = get_noise(
         1,
@@ -108,7 +115,6 @@ def generate_image(
         dtype=torch.bfloat16 if flux_generator.device.type == "cuda" else torch.float32,
         seed=opts.seed,
     )
-    print(x)
     timesteps = get_schedule(
         opts.num_steps,
         x.shape[-1] * x.shape[-2] // 4,
@@ -123,13 +129,11 @@ def generate_image(
     inp = prepare(t5=flux_generator.t5, clip=flux_generator.clip, img=x, prompt=opts.prompt)
     inp_neg = prepare(t5=flux_generator.t5, clip=flux_generator.clip, img=x, prompt=neg_prompt) if use_true_cfg else None
 
-    # offload TEs to CPU, load model to gpu
     if flux_generator.offload:
         flux_generator.t5, flux_generator.clip = flux_generator.t5.cpu(), flux_generator.clip.cpu()
         torch.cuda.empty_cache()
         flux_generator.model = flux_generator.model.to(flux_generator.device)
 
-    # denoise initial noise
     x = denoise(
         flux_generator.model,
         **inp,
@@ -146,13 +150,11 @@ def generate_image(
         neg_vec=inp_neg["vec"] if use_true_cfg else None,
     )
 
-    # offload model, load autoencoder to gpu
     if flux_generator.offload:
         flux_generator.model.cpu()
         torch.cuda.empty_cache()
         flux_generator.ae.decoder.to(x.device)
 
-    # decode latents to pixel space
     x = unpack(x.float(), opts.height, opts.width)
     with torch.autocast(
         device_type=flux_generator.device.type,
@@ -167,13 +169,16 @@ def generate_image(
     t1 = time.perf_counter()
     print(f"Done in {t1 - t0:.1f}s.")
 
-    # bring into PIL format
+    # tensor [-1,1] → uint8 HWC
     x = x.clamp(-1, 1)
     x = rearrange(x[0], "c h w -> h w c")
     img = Image.fromarray((127.5 * (x + 1.0)).cpu().byte().numpy()).convert("RGB")
 
-    # Gallery용 디버그 이미지들을 PIL 리스트로 강제 변환
-    debug = []
+    # 메인 이미지는 파일 경로로 반환 (대용량 base64 전송 이슈 회피)
+    out_path = _save_pil(img, "flux")
+
+    # 디버그 갤러리는 선택적으로 축소/파일 저장
+    debug_paths = []
     for it in (flux_generator.pulid_model.debug_img_list or []):
         try:
             if isinstance(it, Image.Image):
@@ -188,22 +193,41 @@ def generate_image(
                 if arr.dtype != np.uint8:
                     arr = np.clip(arr, 0, 255).astype(np.uint8)
                 pil = Image.fromarray(arr).convert("RGB")
-            debug.append(pil)
+            # 썸네일화 (너비 512)
+            w, h = pil.size
+            if w > 512:
+                nh = int(h * (512 / w))
+                pil = pil.resize((512, nh), Image.BICUBIC)
+            debug_paths.append(_save_pil(pil, "debug"))
         except Exception:
             continue
 
-    return img, str(opts.seed), debug
+    return out_path, str(opts.seed), debug_paths
 
 
 def create_demo(args, model_name: str, device: str = "cuda" if torch.cuda.is_available() else "cpu", offload: bool = False):
-    with gr.Blocks(theme="soft") as demo:
+    custom_css = """
+    /* 상단에 패딩을 강제해 HF 툴바/노치와 겹침 방지 */
+    :root { scroll-padding-top: 72px; }
+    html, body { padding-top: env(safe-area-inset-top); }
+    .gradio-container { padding-top: 16px !important; overflow: visible !important; }
+    /* 상단 배지 영역이 다른 요소 뒤로 깔리지 않도록 */
+    #top-badges { position: relative; z-index: 3; }
+    /* 모바일에서 상단 요소 잘림 방지 */
+    @media (max-width: 640px) {
+      .gradio-container { padding-top: 20px !important; }
+    }
+    """
+
+    with gr.Blocks(theme="soft", css=custom_css) as demo:
+        # 최상단 여백 확보용 스페이서 (툴바가 가리는 환경 대비)
+        gr.HTML("<div style='height: 12px;'></div>")
         gr.HTML(
             """
-            <div class='container' style='display:flex; justify-content:center; gap:12px;'>
+            <div id="top-badges" class='container' style='display:flex; justify-content:center; gap:12px; margin-top:0;'>
                 <a href="https://huggingface.co/spaces/openfree/Best-AI" target="_blank">
                     <img src="https://img.shields.io/static/v1?label=OpenFree&message=BEST%20AI%20Services&color=%230000ff&labelColor=%23000080&logo=huggingface&logoColor=%23ffa500&style=for-the-badge" alt="OpenFree badge">
                 </a>
-
                 <a href="https://discord.gg/openfreeai" target="_blank">
                     <img src="https://img.shields.io/static/v1?label=Discord&message=Openfree%20AI&color=%230000ff&labelColor=%23800080&logo=discord&logoColor=white&style=for-the-badge" alt="Discord badge">
                 </a>
@@ -239,10 +263,11 @@ def create_demo(args, model_name: str, device: str = "cuda" if torch.cuda.is_ava
                 generate_btn = gr.Button("Generate")
 
             with gr.Column():
-                output_image = gr.Image(label="Generated Image", type="pil", format="png")
+                # 파일 경로 모드로 전송 → 브라우저 랜더링 안정적
+                output_image = gr.Image(label="Generated Image", type="filepath", show_download_button=True)
                 seed_output = gr.Textbox(label="Used Seed")
                 intermediate_output = gr.Gallery(
-                    label="Output",
+                    label="Output (dev only)",
                     elem_id="gallery",
                     visible=args.dev,
                     allow_preview=True,
@@ -326,6 +351,6 @@ if __name__ == "__main__":
         huggingface_hub.login(hf_token)
 
     demo = create_demo(args, args.name, args.device, args.offload)
-    # SSR 비활성화: 업로드/탭 이동 중 ClientDisconnect 노이즈 감소 및 렌더 안정화
+    # SSR 비활성화: 렌더 안정화
     demo.launch(ssr_mode=False)