Update app.py

app.py

@@ -1,8 +1,5 @@
 import time
-
 import gradio as gr
-import spaces
-import numpy as np
 import torch
 from einops import rearrange, repeat
 from PIL import Image
@@ -13,12 +10,13 @@ from flux.util import (
     load_ae,
     load_clip,
     load_flow_model,
-    load_flow_model_quintized,
     load_t5,
 )
 from pulid.pipeline_flux import PuLIDPipeline
 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):
     t5 = load_t5(device, max_length=128)
@@ -34,7 +32,7 @@ class FluxGenerator:
         self.device = torch.device('cuda')
         self.offload = False
         self.model_name = 'flux-dev'
-        self.model, self.ae, self.t5, self.clip_model  = get_models(
+        self.model, self.ae, self.t5, self.clip_model = get_models(
             self.model_name,
             device=self.device,
             offload=self.offload,
@@ -46,7 +44,6 @@ class FluxGenerator:
 flux_generator = FluxGenerator()
 
 
-@spaces.GPU(duration=120)
 @torch.inference_mode()
 def generate_image(
     prompt: str,
@@ -71,17 +68,9 @@ def generate_image(
     perform_editing: bool = True,
     inversion_true_cfg: float = 1.0,
 ):
-    """
-    Core function that performs the image generation.
-    """
-    # self.t5.to(self.device)
-    # self.clip_model.to(self.device)
-    # self.ae.to(self.device)
-    # self.model.to(self.device)
-
     flux_generator.t5.max_length = max_sequence_length
 
-    # If seed == -1, random
+    # 시드 설정
     seed = int(seed)
     if seed == -1:
         seed = None
@@ -99,14 +88,12 @@ def generate_image(
         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) Prepare input noise
+    # 1) 입력 노이즈 준비
     noise = get_noise(
         num_samples=1,
         height=opts.height,
@@ -119,65 +106,42 @@ def generate_image(
     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)
-    # Encode id_image directly here
-    encode_t0 = time.perf_counter()
 
-    # Resize image
+    # ID 이미지 인코딩
+    encode_t0 = time.perf_counter()
     id_image = id_image.resize((opts.width, opts.height), resample=Image.LANCZOS)
-
-    # Convert image to torch.Tensor and scale to [-1, 1]
-    x = np.array(id_image).astype(np.float32)
-    x = torch.from_numpy(x)  # shape: (H, W, C)
-    x = (x / 127.5) - 1.0    # now in [-1, 1]
-    x = rearrange(x, "h w c -> 1 c h w")  # shape: (1, C, H, W)
+    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)
-    # Encode with autocast
+    
     with torch.autocast(device_type=flux_generator.device.type, dtype=torch.bfloat16):
         x = flux_generator.ae.encode(x)
-
     x = x.to(torch.bfloat16)
 
-    # Offload if needed
-    if flux_generator.offload:
-        flux_generator.ae.encoder.to("cpu")
-        torch.cuda.empty_cache()
-
     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) Prepare text embeddings
-    if flux_generator.offload:
-        flux_generator.t5 = flux_generator.t5.to(flux_generator.device)
-        flux_generator.clip_model = flux_generator.clip_model.to(flux_generator.device)
-
+    # 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)
 
-    # Offload text encoders, load ID detection to GPU
-    if flux_generator.offload:
-        flux_generator.t5 = flux_generator.t5.cpu()
-        flux_generator.clip_model = flux_generator.clip_model.cpu()
-        torch.cuda.empty_cache()
-        flux_generator.pulid_model.components_to_device(torch.device("cuda"))
-
-    # 3) ID Embeddings (optional)
+    # 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)
-    else:
-        id_embeddings = None
-        uncond_id_embeddings = None
 
     y_0 = inp["img"].clone().detach()
 
+    # 이미지 처리 과정
     inverted = None
     if perform_inversion:
         inverted = rf_inversion(
@@ -198,10 +162,10 @@ def generate_image(
             y_1=noise,
             gamma=gamma
         )
-
         img = inverted
     else:
         img = noise
+        
     inp["img"] = img
     inp_inversion["img"] = img
 
@@ -251,13 +215,7 @@ def generate_image(
             tau=tau,
         )
 
-    # Offload flux model, load auto-decoder
-    if flux_generator.offload:
-        flux_generator.model.cpu()
-        torch.cuda.empty_cache()
-        flux_generator.ae.decoder.to(x.device)
-
-    # 5) Decode latents
+    # 결과 이미지 디코딩
     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):
@@ -273,14 +231,10 @@ def generate_image(
         with torch.autocast(device_type=flux_generator.device.type, dtype=torch.bfloat16):
             recon = flux_generator.ae.decode(recon)
 
-    if flux_generator.offload:
-        flux_generator.ae.decoder.cpu()
-        torch.cuda.empty_cache()
-
     t1 = time.perf_counter()
     print(f"Done in {t1 - t0:.2f} seconds.")
 
-    # Convert to PIL
+    # PIL 이미지로 변환
     if edited is not None:
         edited = edited.clamp(-1, 1)
         edited = rearrange(edited[0], "c h w -> h w c")
@@ -299,84 +253,58 @@ def generate_image(
     return edited, str(opts.seed), flux_generator.pulid_model.debug_img_list
 
 
-
-
-def create_demo(args, model_name: str, device: str = "cuda" if torch.cuda.is_available() else "cpu",
-                offload: bool = False, aggressive_offload: bool = False):
-
-    with gr.Blocks(theme = "apriel") as demo:
-
+def create_demo(args):
+    with gr.Blocks(theme="apriel") as demo:
         with gr.Row():
             with gr.Column():
-                prompt = gr.Textbox(label="Prompt", value="portrait, color, cinematic")
-                id_image = gr.Image(label="ID Image", type="pil")
-                id_weight = gr.Slider(0.0, 1.0, 0.4, step=0.05, label="id weight")
-
-                width = gr.Slider(256, 1536, 1024, step=16, label="Width", visible=args.dev)
-                height = gr.Slider(256, 1536, 1024, step=16, label="Height", visible=args.dev)
-                num_steps = gr.Slider(1, 24, 16, step=1, label="Number of steps")
-                guidance = gr.Slider(1.0, 10.0, 3.5, step=0.1, label="Guidance")
-
-                with gr.Accordion("Advanced Options (True CFG, true_cfg_scale=1 means use fake CFG, >1 means use true CFG", open=False):    # noqa E501
-                    neg_prompt = gr.Textbox(
-                        label="Negative Prompt",
-                        value="")
-                    true_cfg = gr.Slider(1.0, 10.0, 3.5, step=0.1, label="true CFG scale")
-                    timestep_to_start_cfg = gr.Slider(0, 20, 1, step=1, label="timestep to start cfg", visible=args.dev)
-                    start_step = gr.Slider(0, 10, 0, step=1, label="timestep to start inserting ID")
-                    seed = gr.Textbox(-1, label="Seed (-1 for random)")
-                    max_sequence_length = gr.Slider(128, 512, 128, step=128,
-                                                    label="max_sequence_length for prompt (T5), small will be faster")
-                    gr.Markdown("### RF Inversion Options")
-                    gamma = gr.Slider(0.0, 1.0, 0.5, step=0.1, label="gamma")
-                    eta = gr.Slider(0.0, 1.0, 0.8, step=0.1, label="eta")
-                    s = gr.Slider(0.0, 1.0, 0.0, step=0.1, label="s")
-                    tau = gr.Slider(0, 20, 2, step=1, label="tau")
-
-                generate_btn = gr.Button("Generate")
+                prompt = gr.Textbox(label="프롬프트", value="portrait, color, cinematic")
+                id_image = gr.Image(label="ID 이미지", type="pil")
+                id_weight = gr.Slider(0.0, 1.0, 0.4, step=0.05, label="ID 가중치")
+                num_steps = gr.Slider(1, 24, 16, step=1, label="단계 수")
+                guidance = gr.Slider(1.0, 10.0, 3.5, step=0.1, label="가이던스")
+
+                with gr.Accordion("고급 옵션", open=False):
+                    neg_prompt = gr.Textbox(label="네거티브 프롬프트", value="")
+                    true_cfg = gr.Slider(1.0, 10.0, 3.5, step=0.1, label="CFG 스케일")
+                    seed = gr.Textbox(-1, label="시드 (-1: 랜덤)")
+                    gr.Markdown("### 기타 옵션")
+                    gamma = gr.Slider(0.0, 1.0, 0.5, step=0.1, label="감마")
+                    eta = gr.Slider(0.0, 1.0, 0.8, step=0.1, label="에타")
+
+                generate_btn = gr.Button("이미지 생성")
 
             with gr.Column():
-                output_image = gr.Image(label="Generated Image")
-                seed_output = gr.Textbox(label="Used Seed")
-                intermediate_output = gr.Gallery(label='Output', elem_id="gallery", visible=args.dev)
-
-
-        with gr.Row(), gr.Column():
-                gr.Markdown("## Examples")
-                example_inps = [
-                    [
-                        'a portrait of a clown',
-                        'example_inputs/unsplash/lhon-karwan-11tbHtK5STE-unsplash.jpg',
-                        0.5, 3.5, 42, 5.0, 0.7
-                    ],
-                    [
-                        'a portrait of a zombie',
-                        'example_inputs/unsplash/baruk-granda-cfLL_jHQ-Iw-unsplash.jpg',
-                        0.4, 3.5, 42, 5.0, 0.7
-                    ],
-                    [
-                        'a portrait of an elf',
-                        'example_inputs/unsplash/masoud-razeghi--qsrZhXPius-unsplash.jpg',
-                        0.5, 3.5, 42, 5.0, 0.7
-                    ],
-                    [
-                        'a portrait of a demon',
-                        'example_inputs/unsplash/marcin-sajur-nZdMgqvYPBY-unsplash.jpg',
-                        0.3, 3.5, 42, 5.0, 0.7
-                    ],
-                    [
-                        'a portrait of a superhero',
-                        'example_inputs/unsplash/gus-tu-njana-Mf4MN7MZqcE-unsplash.jpg',
-                        0.2, 3.5, 42, 5.0, 0.8
-                    ],
+                output_image = gr.Image(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, true_cfg, eta])
+            ]
+            gr.Examples(examples=example_inps, inputs=[prompt, id_image, id_weight, guidance, seed])
 
         generate_btn.click(
             fn=generate_image,
-            inputs=[prompt, id_image, width, height, num_steps, start_step, guidance, seed, id_weight, neg_prompt,
-                    true_cfg, timestep_to_start_cfg, max_sequence_length, gamma, eta, s, tau],
-            outputs=[output_image, seed_output, intermediate_output],
+            inputs=[prompt, id_image, 512, 512, num_steps, 0, guidance, seed, id_weight, neg_prompt,
+                    true_cfg, 1, 128, gamma, eta, 0, 5],
+            outputs=[output_image, seed_output],
         )
 
     return demo
@@ -384,30 +312,18 @@ def create_demo(args, model_name: str, device: str = "cuda" if torch.cuda.is_ava
 
 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', help='version of the model', choices=['v0.9.0', 'v0.9.1'])
-    parser.add_argument("--name", type=str, default="flux-dev", choices=list('flux-dev'),
-                        help="currently only support flux-dev")
-    parser.add_argument("--device", type=str, default="cuda", help="Device to use")
-    parser.add_argument("--offload", action="store_true", help="Offload model to CPU when not in use")
-    parser.add_argument("--aggressive_offload", action="store_true", help="Offload model more aggressively to CPU when not in use, for 24G GPUs")
-    parser.add_argument("--fp8", action="store_true", help="use flux-dev-fp8 model")
-    parser.add_argument("--onnx_provider", type=str, default="gpu", choices=["gpu", "cpu"],
-                        help="set onnx_provider to cpu (default gpu) can help reduce RAM usage, and when combined with"
-                             "fp8 option, the peak RAM is under 15GB")
-    parser.add_argument("--port", type=int, default=8080, help="Port to use")
-    parser.add_argument("--dev", action='store_true', help="Development mode")
-    parser.add_argument("--pretrained_model", type=str, help='for development')
+    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()
 
-    # args.fp8 = True
-    if args.aggressive_offload:
-        args.offload = True
-
     print(f"Using device: {args.device}")
-    print(f"fp8: {args.fp8}")
     print(f"Offload: {args.offload}")
 
-    demo = create_demo(args, args.name, args.device, args.offload, args.aggressive_offload)
-    demo.launch(ssr_mode=False)
+    demo = create_demo(args)
+    demo.launch(ssr_mode=False)