modified: gradio_app_asy.py

gradio_app_asy.py

@@ -64,6 +64,7 @@ def download_file(repo_id, file_name):
 # Load model function with dynamic paths based on the selected model
 def load_target_model(frame, domain):
     global model, clip_l, t5xxl, ae, lora_model
+    
     BASE_FLUX_CHECKPOINT=download_file(flux_repo_id, flux_file)
     CLIP_L_PATH = download_file(clip_repo_id, clip_l_file)
     T5XXL_PATH = download_file(clip_repo_id, t5xxl_file)
@@ -71,17 +72,15 @@ def load_target_model(frame, domain):
     LORA_WEIGHTS_PATH = download_file(lora_repo_id, lora_paths[frame])
 
     logger.info("Loading models...")
-    # try:
-    if model is None is None or clip_l is None or t5xxl is None or ae is None:
-        _, model = flux_utils.load_flow_model(
-            BASE_FLUX_CHECKPOINT, torch.float8_e4m3fn, "cpu", disable_mmap=False
-        )
-        clip_l = flux_utils.load_clip_l(CLIP_L_PATH, torch.bfloat16, "cpu", disable_mmap=False)
-        clip_l.eval()
-        t5xxl = flux_utils.load_t5xxl(T5XXL_PATH, torch.bfloat16, "cpu", disable_mmap=False)
-        t5xxl.eval()
-        ae = flux_utils.load_ae(AE_PATH, torch.bfloat16, "cpu", disable_mmap=False)
-        logger.info("Models loaded successfully.")
+    _, model = flux_utils.load_flow_model(
+        BASE_FLUX_CHECKPOINT, torch.float8_e4m3fn, "cpu", disable_mmap=False
+    )
+    clip_l = flux_utils.load_clip_l(CLIP_L_PATH, torch.bfloat16, "cpu", disable_mmap=False)
+    clip_l.eval()
+    t5xxl = flux_utils.load_t5xxl(T5XXL_PATH, torch.bfloat16, "cpu", disable_mmap=False)
+    t5xxl.eval()
+    ae = flux_utils.load_ae(AE_PATH, torch.bfloat16, "cpu", disable_mmap=False)
+    logger.info("Models loaded successfully.")
     # Load LoRA weights
     multiplier = 1.0
     weights_sd = load_file(LORA_WEIGHTS_PATH)
@@ -197,6 +196,8 @@ def infer(prompt, frame, seed=0):
     generated_image = Image.fromarray((127.5 * (x + 1.0)).float().cpu().numpy().astype(np.uint8)[0])
 
     logger.info("Image generation completed.")
+    torch.cuda.empty_cache()
+
     return generated_image
 
 def update_domains(floor):
@@ -232,7 +233,7 @@ with gr.Blocks() as demo:
         with gr.Column(scale=1):
             # Status message box
             status_box = gr.Textbox(label="Status", placeholder="Model loading status", interactive=False, value="Model not loaded", lines=3)
-            
+
     with gr.Row():
         with gr.Column(scale=1):
             # Input for the prompt