Update app.py

app.py

@@ -3,6 +3,7 @@ import gradio as gr
 import torch
 from einops import rearrange, repeat
 from PIL import Image
+import numpy as np
 
 from flux.sampling import denoise, get_noise, get_schedule, prepare, rf_denoise, rf_inversion, unpack
 from flux.util import (
@@ -18,7 +19,16 @@ from pulid.utils import resize_numpy_image_long, seed_everything
 # 간단한 인용 정보 추가
 _CITE_ = """PuLID: Person-under-Language Image Diffusion Model"""
 
-def get_models(name: str, device: torch.device, offload: bool):
+# GPU 사용 가능 여부 확인 및 장치 설정
+def get_device():
+    if torch.cuda.is_available():
+        return torch.device('cuda')
+    else:
+        print("CUDA GPU를 찾을 수 없습니다. CPU를 사용합니다.")
+        return torch.device('cpu')
+
+def get_models(name: str, device, offload: bool):
+    print(f"모델을 {device}에 로드합니다.")
     t5 = load_t5(device, max_length=128)
     clip_model = load_clip(device)
     model = load_flow_model(name, device="cpu" if offload else device)
@@ -29,19 +39,36 @@ def get_models(name: str, device: torch.device, offload: bool):
 
 class FluxGenerator:
     def __init__(self):
-        self.device = torch.device('cuda')
+        # GPU 사용 가능 여부에 따라 장치 설정
+        self.device = get_device()
         self.offload = False
         self.model_name = 'flux-dev'
-        self.model, self.ae, self.t5, self.clip_model = get_models(
-            self.model_name,
-            device=self.device,
-            offload=self.offload,
-        )
-        self.pulid_model = PuLIDPipeline(self.model, 'cuda', weight_dtype=torch.bfloat16)
-        self.pulid_model.load_pretrain()
-
-
-flux_generator = FluxGenerator()
+        
+        # 모델 로드 시도
+        try:
+            self.model, self.ae, self.t5, self.clip_model = get_models(
+                self.model_name,
+                device=self.device,
+                offload=self.offload,
+            )
+            self.pulid_model = PuLIDPipeline(
+                self.model, 
+                'cuda' if torch.cuda.is_available() else 'cpu', 
+                weight_dtype=torch.bfloat16 if self.device.type == 'cuda' else torch.float32
+            )
+            self.pulid_model.load_pretrain()
+            self.initialized = True
+        except Exception as e:
+            print(f"모델 초기화 중 오류 발생: {e}")
+            self.initialized = False
+
+# 모델 초기화 시도
+try:
+    flux_generator = FluxGenerator()
+    model_initialized = flux_generator.initialized
+except Exception as e:
+    print(f"FluxGenerator 초기화 중 오류 발생: {e}")
+    model_initialized = False
 
 
 @torch.inference_mode()
@@ -63,87 +90,91 @@ def generate_image(
     eta: float = 0.7,
     s: float = 0,
     tau: float = 5,
-    perform_inversion: bool = True,
-    perform_reconstruction: bool = False,
-    perform_editing: bool = True,
-    inversion_true_cfg: float = 1.0,
 ):
-    flux_generator.t5.max_length = max_sequence_length
-
-    # 시드 설정
-    seed = int(seed)
-    if seed == -1:
-        seed = None
-
-    opts = SamplingOptions(
-        prompt=prompt,
-        width=width,
-        height=height,
-        num_steps=num_steps,
-        guidance=guidance,
-        seed=seed,
-    )
-
-    if opts.seed is None:
-        opts.seed = torch.Generator(device="cpu").seed()
-
-    seed_everything(opts.seed)
-    print(f"Generating prompt: '{opts.prompt}' (seed={opts.seed})...")
-    t0 = time.perf_counter()
-
-    use_true_cfg = abs(true_cfg - 1.0) > 1e-6
-
-    # 1) 입력 노이즈 준비
-    noise = get_noise(
-        num_samples=1,
-        height=opts.height,
-        width=opts.width,
-        device=flux_generator.device,
-        dtype=torch.bfloat16,
-        seed=opts.seed,
-    )
-    bs, c, h, w = noise.shape
-    noise = rearrange(noise, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
-    if noise.shape[0] == 1 and bs > 1:
-        noise = repeat(noise, "1 ... -> bs ...", bs=bs)
-
-    # ID 이미지 인코딩
-    encode_t0 = time.perf_counter()
-    id_image = id_image.resize((opts.width, opts.height), resample=Image.LANCZOS)
-    x = torch.from_numpy(np.array(id_image).astype(np.float32))
-    x = (x / 127.5) - 1.0
-    x = rearrange(x, "h w c -> 1 c h w")
-    x = x.to(flux_generator.device)
+    # 모델이 초기화되지 않았으면 오류 메시지 반환
+    if not model_initialized:
+        return None, "GPU 오류: CUDA GPU를 찾을 수 없어 모델을 초기화할 수 없습니다.", None
     
-    with torch.autocast(device_type=flux_generator.device.type, dtype=torch.bfloat16):
-        x = flux_generator.ae.encode(x)
-    x = x.to(torch.bfloat16)
-
-    encode_t1 = time.perf_counter()
-    print(f"Encoded in {encode_t1 - encode_t0:.2f} seconds.")
+    # ID 이미지가 없으면 실행 불가
+    if id_image is None:
+        return None, "오류: ID 이미지가 필요합니다.", None
     
-    timesteps = get_schedule(opts.num_steps, x.shape[-1] * x.shape[-2] // 4, shift=False)
-
-    # 2) 텍스트 임베딩 준비
-    inp = prepare(t5=flux_generator.t5, clip=flux_generator.clip_model, img=x, prompt=opts.prompt)
-    inp_inversion = prepare(t5=flux_generator.t5, clip=flux_generator.clip_model, img=x, prompt="")
-    inp_neg = None
-    if use_true_cfg:
-        inp_neg = prepare(t5=flux_generator.t5, clip=flux_generator.clip_model, img=x, prompt=neg_prompt)
-
-    # 3) ID 임베딩 생성
-    id_embeddings = None
-    uncond_id_embeddings = None
-    if id_image is not None:
-        id_image = np.array(id_image)
-        id_image = resize_numpy_image_long(id_image, 1024)
-        id_embeddings, uncond_id_embeddings = flux_generator.pulid_model.get_id_embedding(id_image, cal_uncond=use_true_cfg)
-
-    y_0 = inp["img"].clone().detach()
-
-    # 이미지 처리 과정
-    inverted = None
-    if perform_inversion:
+    try:
+        flux_generator.t5.max_length = max_sequence_length
+
+        # 시드 설정
+        seed = int(seed)
+        if seed == -1:
+            seed = None
+
+        opts = SamplingOptions(
+            prompt=prompt,
+            width=width,
+            height=height,
+            num_steps=num_steps,
+            guidance=guidance,
+            seed=seed,
+        )
+
+        if opts.seed is None:
+            opts.seed = torch.Generator(device="cpu").seed()
+
+        seed_everything(opts.seed)
+        print(f"Generating prompt: '{opts.prompt}' (seed={opts.seed})...")
+        t0 = time.perf_counter()
+
+        use_true_cfg = abs(true_cfg - 1.0) > 1e-6
+
+        # 1) 입력 노이즈 준비
+        noise = get_noise(
+            num_samples=1,
+            height=opts.height,
+            width=opts.width,
+            device=flux_generator.device,
+            dtype=torch.bfloat16 if flux_generator.device.type == 'cuda' else torch.float32,
+            seed=opts.seed,
+        )
+        bs, c, h, w = noise.shape
+        noise = rearrange(noise, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
+        if noise.shape[0] == 1 and bs > 1:
+            noise = repeat(noise, "1 ... -> bs ...", bs=bs)
+
+        # ID 이미지 인코딩
+        encode_t0 = time.perf_counter()
+        id_image = id_image.resize((opts.width, opts.height), resample=Image.LANCZOS)
+        x = torch.from_numpy(np.array(id_image).astype(np.float32))
+        x = (x / 127.5) - 1.0
+        x = rearrange(x, "h w c -> 1 c h w")
+        x = x.to(flux_generator.device)
+        
+        dtype = torch.bfloat16 if flux_generator.device.type == 'cuda' else torch.float32
+        with torch.autocast(device_type=flux_generator.device.type, dtype=dtype):
+            x = flux_generator.ae.encode(x)
+        x = x.to(dtype)
+
+        encode_t1 = time.perf_counter()
+        print(f"Encoded in {encode_t1 - encode_t0:.2f} seconds.")
+        
+        timesteps = get_schedule(opts.num_steps, x.shape[-1] * x.shape[-2] // 4, shift=False)
+
+        # 2) 텍스트 임베딩 준비
+        inp = prepare(t5=flux_generator.t5, clip=flux_generator.clip_model, img=x, prompt=opts.prompt)
+        inp_inversion = prepare(t5=flux_generator.t5, clip=flux_generator.clip_model, img=x, prompt="")
+        inp_neg = None
+        if use_true_cfg:
+            inp_neg = prepare(t5=flux_generator.t5, clip=flux_generator.clip_model, img=x, prompt=neg_prompt)
+
+        # 3) ID 임베딩 생성
+        id_embeddings = None
+        uncond_id_embeddings = None
+        if id_image is not None:
+            id_image = np.array(id_image)
+            id_image = resize_numpy_image_long(id_image, 1024)
+            id_embeddings, uncond_id_embeddings = flux_generator.pulid_model.get_id_embedding(id_image, cal_uncond=use_true_cfg)
+
+        y_0 = inp["img"].clone().detach()
+
+        # 이미지 처리 과정
         inverted = rf_inversion(
             flux_generator.model,
             **inp_inversion,
@@ -153,7 +184,7 @@ def generate_image(
             id_weight=id_weight,
             start_step=start_step,
             uncond_id=uncond_id_embeddings,
-            true_cfg=inversion_true_cfg,
+            true_cfg=true_cfg,
             timestep_to_start_cfg=timestep_to_start_cfg,
             neg_txt=inp_neg["txt"] if use_true_cfg else None,
             neg_txt_ids=inp_neg["txt_ids"] if use_true_cfg else None,
@@ -162,38 +193,10 @@ def generate_image(
             y_1=noise,
             gamma=gamma
         )
-        img = inverted
-    else:
-        img = noise
         
-    inp["img"] = img
-    inp_inversion["img"] = img
-
-    recon = None
-    if perform_reconstruction:
-        recon = rf_denoise(
-            flux_generator.model,
-            **inp_inversion,
-            timesteps=timesteps,
-            guidance=opts.guidance,
-            id=id_embeddings,
-            id_weight=id_weight,
-            start_step=start_step,
-            uncond_id=uncond_id_embeddings,
-            true_cfg=inversion_true_cfg,
-            timestep_to_start_cfg=timestep_to_start_cfg,
-            neg_txt=inp_neg["txt"] if use_true_cfg else None,
-            neg_txt_ids=inp_neg["txt_ids"] if use_true_cfg else None,
-            neg_vec=inp_neg["vec"] if use_true_cfg else None,
-            aggressive_offload=False,
-            y_0=y_0,
-            eta=eta,
-            s=s,
-            tau=tau,
-        )
+        inp["img"] = inverted
+        inp_inversion["img"] = inverted
 
-    edited = None
-    if perform_editing:
         edited = rf_denoise(
             flux_generator.model,
             **inp,
@@ -215,46 +218,37 @@ def generate_image(
             tau=tau,
         )
 
-    # 결과 이미지 디코딩
-    if edited is not None:
+        # 결과 이미지 디코딩
         edited = unpack(edited.float(), opts.height, opts.width)
-        with torch.autocast(device_type=flux_generator.device.type, dtype=torch.bfloat16):
+        with torch.autocast(device_type=flux_generator.device.type, dtype=dtype):
             edited = flux_generator.ae.decode(edited)
 
-    if inverted is not None:
-        inverted = unpack(inverted.float(), opts.height, opts.width)
-        with torch.autocast(device_type=flux_generator.device.type, dtype=torch.bfloat16):
-            inverted = flux_generator.ae.decode(inverted)
-    
-    if recon is not None:
-        recon = unpack(recon.float(), opts.height, opts.width)
-        with torch.autocast(device_type=flux_generator.device.type, dtype=torch.bfloat16):
-            recon = flux_generator.ae.decode(recon)
-
-    t1 = time.perf_counter()
-    print(f"Done in {t1 - t0:.2f} seconds.")
+        t1 = time.perf_counter()
+        print(f"Done in {t1 - t0:.2f} seconds.")
 
-    # PIL 이미지로 변환
-    if edited is not None:
+        # PIL 이미지로 변환
         edited = edited.clamp(-1, 1)
         edited = rearrange(edited[0], "c h w -> h w c")
         edited = Image.fromarray((127.5 * (edited + 1.0)).cpu().byte().numpy())
 
-    if inverted is not None:
-        inverted = inverted.clamp(-1, 1)
-        inverted = rearrange(inverted[0], "c h w -> h w c")
-        inverted = Image.fromarray((127.5 * (inverted + 1.0)).cpu().byte().numpy())
+        return edited, str(opts.seed), flux_generator.pulid_model.debug_img_list
     
-    if recon is not None:
-        recon = recon.clamp(-1, 1)
-        recon = rearrange(recon[0], "c h w -> h w c")
-        recon = Image.fromarray((127.5 * (recon + 1.0)).cpu().byte().numpy())
+    except Exception as e:
+        import traceback
+        error_msg = f"이미지 생성 중 오류 발생: {str(e)}\n{traceback.format_exc()}"
+        print(error_msg)
+        return None, error_msg, None
 
-    return edited, str(opts.seed), flux_generator.pulid_model.debug_img_list
 
-
-def create_demo(args):
-    with gr.Blocks(theme="apriel") as demo:
+def create_demo():
+    with gr.Blocks() as demo:
+        gr.Markdown("# PuLID: 인물 이미지 변환 도구")
+        
+        if not model_initialized:
+            gr.Markdown("## ⚠️ 오류: CUDA GPU를 찾을 수 없습니다")
+            gr.Markdown("이 응용 프로그램은 CUDA 지원 GPU가 필요합니다. CPU에서는 실행할 수 없습니다.")
+            return demo
+            
         with gr.Row():
             with gr.Column():
                 prompt = gr.Textbox(label="프롬프트", value="portrait, color, cinematic")
@@ -275,31 +269,9 @@ def create_demo(args):
 
             with gr.Column():
                 output_image = gr.Image(label="생성된 이미지")
-                seed_output = gr.Textbox(label="사용된 시드")
+                seed_output = gr.Textbox(label="결과/오류 메시지")
                 gr.Markdown(_CITE_)
 
-        # 예제 입력
-        with gr.Row():
-            gr.Markdown("## 예제")
-            example_inps = [
-                [
-                    'a portrait of a clown',
-                    'example_inputs/unsplash/lhon-karwan-11tbHtK5STE-unsplash.jpg',
-                    0.5, 3.5, 42
-                ],
-                [
-                    'a portrait of a zombie',
-                    'example_inputs/unsplash/baruk-granda-cfLL_jHQ-Iw-unsplash.jpg',
-                    0.4, 3.5, 42
-                ],
-                [
-                    'a portrait of an elf',
-                    'example_inputs/unsplash/masoud-razeghi--qsrZhXPius-unsplash.jpg',
-                    0.5, 3.5, 42
-                ]
-            ]
-            gr.Examples(examples=example_inps, inputs=[prompt, id_image, id_weight, guidance, seed])
-
         generate_btn.click(
             fn=generate_image,
             inputs=[prompt, id_image, 512, 512, num_steps, 0, guidance, seed, id_weight, neg_prompt,
@@ -312,18 +284,12 @@ def create_demo(args):
 
 if __name__ == "__main__":
     import argparse
-    import numpy as np
 
     parser = argparse.ArgumentParser(description="PuLID for FLUX.1-dev")
     parser.add_argument('--version', type=str, default='v0.9.1')
     parser.add_argument("--name", type=str, default="flux-dev")
-    parser.add_argument("--device", type=str, default="cuda")
-    parser.add_argument("--offload", action="store_true")
     parser.add_argument("--port", type=int, default=8080)
     args = parser.parse_args()
 
-    print(f"Using device: {args.device}")
-    print(f"Offload: {args.offload}")
-
-    demo = create_demo(args)
+    demo = create_demo()
     demo.launch(ssr_mode=False)