Update app.py

app.py

@@ -12,20 +12,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import gc
+import shutil
 import os
 
 import gradio as gr
-import huggingface_hub
 import pillow_avif
-import spaces
 import torch
-import gc
 from huggingface_hub import snapshot_download
 from pillow_heif import register_heif_opener
-from PIL import Image, ImageDraw, ImageFont
 
 from pipelines.pipeline_infu_flux import InfUFluxPipeline
 
+
 # Register HEIF support for Pillow
 register_heif_opener()
 
@@ -47,9 +46,26 @@ loaded_pipeline_config = {
 
 
 def download_models():
-    snapshot_download(repo_id='ByteDance/InfiniteYou', revision="61248f501725e31c16a90a1c7ffe63f6e2839c65", local_dir='./models/InfiniteYou', local_dir_use_symlinks=False)
+    # InfiniteYou モデル (必要な部分だけ)
+    snapshot_download(
+        repo_id='ByteDance/InfiniteYou',
+        local_dir='./models/InfiniteYou',
+        local_dir_use_symlinks=False,
+        allow_patterns=[
+            "infu_flux_v1.0/aes_stage2/**",   # aes_stage2 モデル
+            "supports/insightface/**",        # 顔認識に必要
+            "supports/optional_loras/**",     # LoRA が必要なら残す
+        ]
+    )
+    
+    # FLUX.1-dev (ベースモデルは必須なので別途取得)
     try:
-        snapshot_download(repo_id='black-forest-labs/FLUX.1-dev', local_dir='./models/FLUX.1-dev', local_dir_use_symlinks=False)
+        snapshot_download(
+            repo_id='black-forest-labs/FLUX.1-dev',
+            local_dir='./models/FLUX.1-dev',
+            local_dir_use_symlinks=False,
+            allow_patterns=["*.safetensors", "*.json", "*.txt"]
+        )
     except Exception as e:
         print(e)
         print('\nYou are downloading `black-forest-labs/FLUX.1-dev` to `./models/FLUX.1-dev` but failed. '
@@ -61,36 +77,11 @@ def download_models():
         exit()
 
 
-def init_pipeline(model_version, enable_realism, enable_anti_blur):
-    loaded_pipeline_config["enable_realism"] = enable_realism
-    loaded_pipeline_config["enable_anti_blur"] = enable_anti_blur
-    loaded_pipeline_config["model_version"] = model_version
-
-    pipeline = loaded_pipeline_config['pipeline']
-    gc.collect()
-    torch.cuda.empty_cache()
-
-    model_path = f'./models/InfiniteYou/infu_flux_v1.0/{model_version}'
-    print(f'loading model from {model_path}')
-
-    pipeline = InfUFluxPipeline(
-        base_model_path='./models/FLUX.1-dev',
-        infu_model_path=model_path,
-        insightface_root_path='./models/InfiniteYou/supports/insightface',
-        image_proj_num_tokens=8,
-        infu_flux_version='v1.0',
-        model_version=model_version,
-    )
-
-    loaded_pipeline_config['pipeline'] = pipeline
-
-    pipeline.pipe.delete_adapters(['realism', 'anti_blur'])
-    loras = []
-    if enable_realism: loras.append(['realism', 1.0])
-    if enable_anti_blur: loras.append(['anti_blur', 1.0])
-    pipeline.load_loras_state_dict(loras)
-
-    return pipeline
+def clean_hf_cache():
+    cache_dir = os.path.expanduser("~/.cache/huggingface")
+    if os.path.exists(cache_dir):
+        print(f"Cleaning Hugging Face cache at {cache_dir}")
+        shutil.rmtree(cache_dir, ignore_errors=True)
 
 
 def prepare_pipeline(model_version, enable_realism, enable_anti_blur):
@@ -107,62 +98,43 @@ def prepare_pipeline(model_version, enable_realism, enable_anti_blur):
     loaded_pipeline_config["model_version"] = model_version
 
     pipeline = loaded_pipeline_config['pipeline']
-    if pipeline is None or pipeline.model_version != model_version: 
+    if pipeline is None or pipeline.model_version != model_version:
         print(f'Switching model to {model_version}')
-        pipeline.model_version = model_version
+        del pipeline
+        del loaded_pipeline_config['pipeline']
+        gc.collect()
+        torch.cuda.empty_cache()
+
         if model_version == 'aes_stage2':
-            pipeline.infusenet_sim.cpu()
-            pipeline.image_proj_model_sim.cpu()
-            torch.cuda.empty_cache()
-            pipeline.infusenet_aes.to('cuda')
-            pipeline.pipe.controlnet = pipeline.infusenet_aes
-            pipeline.image_proj_model_aes.to('cuda')
-            pipeline.image_proj_model = pipeline.image_proj_model_aes
+            model_path = f'./models/InfiniteYou/infu_flux_v1.0/aes_stage2'
+        elif model_version == 'sim_stage1':
+            model_path = f'./models/InfiniteYou/infu_flux_v1.0/sim_stage1'
         else:
-            pipeline.infusenet_aes.cpu()
-            pipeline.image_proj_model_aes.cpu()
-            torch.cuda.empty_cache()
-            pipeline.infusenet_sim.to('cuda')
-            pipeline.pipe.controlnet = pipeline.infusenet_sim
-            pipeline.image_proj_model_sim.to('cuda')
-            pipeline.image_proj_model = pipeline.image_proj_model_sim
+            raise ValueError(f'Model version {model_version} not supported.')
+        print(f'Loading model from {model_path}')
+
+        pipeline = InfUFluxPipeline(
+            base_model_path='./models/FLUX.1-dev',
+            infu_model_path=model_path,
+            insightface_root_path='./models/InfiniteYou/supports/insightface',
+            image_proj_num_tokens=8,
+            infu_flux_version='v1.0',
+            model_version=model_version,
+        )
 
         loaded_pipeline_config['pipeline'] = pipeline
 
     pipeline.pipe.delete_adapters(['realism', 'anti_blur'])
     loras = []
-    if enable_realism: loras.append(['realism', 1.0])
-    if enable_anti_blur: loras.append(['anti_blur', 1.0])
-    pipeline.load_loras_state_dict(loras)
+    if enable_realism:
+        loras.append(['./models/InfiniteYou/supports/optional_loras/flux_realism_lora.safetensors', 'realism', 1.0])
+    if enable_anti_blur:
+        loras.append(['./models/InfiniteYou/supports/optional_loras/flux_anti_blur_lora.safetensors', 'anti_blur', 1.0])
+    pipeline.load_loras(loras)
 
     return pipeline
 
 
-def add_safety_watermark(image, text='AI Generated', font_path=None):
-    width, height = image.size
-    draw = ImageDraw.Draw(image)
-
-    font_size = int(height * 0.028)
-    if font_path:
-        font = ImageFont.truetype(font_path, font_size)
-    else:
-        font = ImageFont.load_default(size=font_size)
-    
-    text_bbox = draw.textbbox((0, 0), text, font=font)
-    text_width, text_height = text_bbox[2] - text_bbox[0], text_bbox[3] - text_bbox[1]
-    x = width - text_width - 10
-    y = height - text_height - 20
-
-    shadow_offset = 2
-    shadow_color = "black"
-    draw.text((x + shadow_offset, y + shadow_offset), text, font=font, fill=shadow_color)
-
-    draw.text((x, y), text, font=font, fill="white")
-
-    return image
-
-
-@spaces.GPU(duration=120)
 def generate_image(
     input_image, 
     control_image, 
@@ -179,12 +151,12 @@ def generate_image(
     enable_anti_blur,
     model_version
 ):
-    try:
-        pipeline = prepare_pipeline(model_version=model_version, enable_realism=enable_realism, enable_anti_blur=enable_anti_blur)
+    pipeline = prepare_pipeline(model_version=model_version, enable_realism=enable_realism, enable_anti_blur=enable_anti_blur)
 
-        if seed == 0:
-            seed = torch.seed() & 0xFFFFFFFF
+    if seed == 0:
+        seed = torch.seed() & 0xFFFFFFFF
 
+    try:
         image = pipeline(
             id_image=input_image,
             prompt=prompt,
@@ -198,13 +170,12 @@ def generate_image(
             infusenet_guidance_start=infusenet_guidance_start,
             infusenet_guidance_end=infusenet_guidance_end,
         )
-        image = add_safety_watermark(image)
     except Exception as e:
         print(e)
         gr.Error(f"An error occurred: {e}")
         return gr.update()
 
-    return gr.update(value=image, label=f"Generated Image, seed = {seed}")
+    return gr.update(value = image, label=f"Generated Image, seed = {seed}")
 
 
 def generate_examples(id_image, control_image, prompt_text, seed, enable_realism, enable_anti_blur, model_version):
@@ -233,8 +204,7 @@ with gr.Blocks() as demo:
         <a href="https://bytedance.github.io/InfiniteYou">[Project Page]</a>&ensp;
         <a href="https://arxiv.org/abs/2503.16418">[Paper]</a>&ensp;
         <a href="https://github.com/bytedance/InfiniteYou">[Code]</a>&ensp;
-        <a href="https://huggingface.co/ByteDance/InfiniteYou">[Model]</a>&ensp;
-        <a href="https://github.com/bytedance/ComfyUI_InfiniteYou">[ComfyUI]</a>
+        <a href="https://huggingface.co/ByteDance/InfiniteYou">[Model]</a>
     </div>
     """)
 
@@ -298,7 +268,7 @@ with gr.Blocks() as demo:
         inputs=[ui_id_image, ui_control_image, ui_prompt_text, ui_seed, ui_enable_realism, ui_enable_anti_blur, ui_model_version],
         outputs=[image_output],
         fn=generate_examples,
-        cache_examples=False
+        cache_examples=True,
     )
 
     ui_btn_generate.click(
@@ -335,12 +305,12 @@ with gr.Blocks() as demo:
         The images used in this demo are sourced from consented subjects or generated by the models. These pictures are intended solely to show the capabilities of our research. If you have any concerns, please contact us, and we will promptly remove any inappropriate content.
         
         The use of the released code, model, and demo must strictly adhere to the respective licenses. 
-        Our code is released under the [Apache License 2.0](https://github.com/bytedance/InfiniteYou/blob/main/LICENSE), 
+        Our code is released under the [Apache 2.0 License](https://github.com/bytedance/InfiniteYou/blob/main/LICENSE), 
         and our model is released under the [Creative Commons Attribution-NonCommercial 4.0 International Public License](https://huggingface.co/ByteDance/InfiniteYou/blob/main/LICENSE) 
         for academic research purposes only. Any manual or automatic downloading of the face models from [InsightFace](https://github.com/deepinsight/insightface), 
         the [FLUX.1-dev](https://huggingface.co/black-forest-labs/FLUX.1-dev) base model, LoRAs, *etc.*, must follow their original licenses and be used only for academic research purposes.
 
-        This research aims to positively impact the field of Generative AI. Any usage of this method must be responsible and comply with local laws. The developers do not assume any responsibility for any potential misuse. We added the "AI Generated" watermark for enhanced safety.
+        This research aims to positively impact the field of Generative AI. Any usage of this method must be responsible and comply with local laws. The developers do not assume any responsibility for any potential misuse.
         """
     )    
 
@@ -363,14 +333,13 @@ with gr.Blocks() as demo:
         """
     )
 
-
-huggingface_hub.login(os.getenv('PRIVATE_HF_TOKEN'))
-
 download_models()
 
-init_pipeline(model_version=ModelVersion.DEFAULT_VERSION, enable_realism=ENABLE_REALISM_DEFAULT, enable_anti_blur=ENABLE_ANTI_BLUR_DEFAULT)
+prepare_pipeline(model_version=ModelVersion.DEFAULT_VERSION, enable_realism=ENABLE_REALISM_DEFAULT, enable_anti_blur=ENABLE_ANTI_BLUR_DEFAULT)
 
-# demo.queue()
-demo.launch()
+demo.queue()
+demo.launch(server_name='localhost')  # localhost
 # demo.launch(server_name='0.0.0.0')  # IPv4
 # demo.launch(server_name='[::]')  # IPv6
+
+clean_hf_cache()