update app.py

Author: EndlessSora

app.py

@@ -79,7 +79,12 @@ def prepare_pipeline(model_version, enable_realism, enable_anti_blur):
         gc.collect()
         torch.cuda.empty_cache()
 
-        model_path = f'./models/InfiniteYou/infu_flux_v1.0/{model_version}'
+        if model_version == 'aes_stage2':
+            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:
+            raise ValueError(f'Model version {model_version} not supported.')
         print(f'Loading model from {model_path}')
 
         pipeline = InfUFluxPipeline(
@@ -307,5 +312,6 @@ def generate_examples(id_image, control_image, prompt_text, seed, enable_realism
 prepare_pipeline(model_version=ModelVersion.DEFAULT_VERSION, enable_realism=ENABLE_REALISM_DEFAULT, enable_anti_blur=ENABLE_ANTI_BLUR_DEFAULT)
 
 demo.queue()
-demo.launch(server_name='0.0.0.0')  # IPv4
+demo.launch(server_name='localhost')  # localhost
+# demo.launch(server_name='0.0.0.0')  # IPv4
 # demo.launch(server_name='[::]')  # IPv6

pipelines/pipeline_infu_flux.py

@@ -194,7 +194,7 @@ def __init__(
             ff_mult=4,
         )
         image_proj_model_path = os.path.join(infu_model_path, 'image_proj_model.bin')
-        ipm_state_dict = torch.load(image_proj_model_path, map_location="cpu")
+        ipm_state_dict = torch.load(image_proj_model_path, map_location="cpu", weights_only=True)
         image_proj_model.load_state_dict(ipm_state_dict['image_proj'])
         del ipm_state_dict
         image_proj_model.to('cuda', torch.bfloat16)