modified: gradio_app.py

gradio_app.py

@@ -94,33 +94,35 @@ def load_target_model(selected_model):
     AE_PATH = download_file(ae_repo_id, ae_file)
     LORA_WEIGHTS_PATH = download_file(lora_repo, lora_file)
 
-    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)
-
-        # Load LoRA weights
-        multiplier = 1.0
-        weights_sd = load_file(LORA_WEIGHTS_PATH)
-        lora_model, _ = lora_flux.create_network_from_weights(multiplier, None, ae, [clip_l, t5xxl], model, weights_sd, True)
-        lora_model.apply_to([clip_l, t5xxl], model)
-        info = lora_model.load_state_dict(weights_sd, strict=True)
-        logger.info(f"Loaded LoRA weights from {LORA_WEIGHTS_PATH}: {info}")
-        lora_model.eval()
-
-        logger.info("Models loaded successfully.")
-        return "Models loaded successfully. Using Recraft: {}".format(selected_model)
-
-    except Exception as e:
-        logger.error(f"Error loading models: {e}")
-        return f"Error loading models: {e}"
+    return "Models loaded successfully. Using Recraft: {}".format(selected_model)
+
+    # 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)
+
+    #     # Load LoRA weights
+    #     multiplier = 1.0
+    #     weights_sd = load_file(LORA_WEIGHTS_PATH)
+    #     lora_model, _ = lora_flux.create_network_from_weights(multiplier, None, ae, [clip_l, t5xxl], model, weights_sd, True)
+    #     lora_model.apply_to([clip_l, t5xxl], model)
+    #     info = lora_model.load_state_dict(weights_sd, strict=True)
+    #     logger.info(f"Loaded LoRA weights from {LORA_WEIGHTS_PATH}: {info}")
+    #     lora_model.eval()
+
+    #     logger.info("Models loaded successfully.")
+    #     return "Models loaded successfully. Using Recraft: {}".format(selected_model)
+
+    # except Exception as e:
+    #     logger.error(f"Error loading models: {e}")
+    #     return f"Error loading models: {e}"
 
 # Image pre-processing (resize and padding)
 class ResizeWithPadding:
@@ -154,10 +156,37 @@ class ResizeWithPadding:
 # The function to generate image from a prompt and conditional image
 @spaces.GPU(duration=180)
 def infer(prompt, sample_image, recraft_model, seed=0):
-    global model, clip_l, t5xxl, ae, lora_model
-    if model is None or lora_model is None or clip_l is None or t5xxl is None or ae is None:
-        logger.error("Models not loaded. Please load the models first.")
-        return None
+    # global model, clip_l, t5xxl, ae, lora_model
+    # if model is None or lora_model is None or clip_l is None or t5xxl is None or ae is None:
+    #     logger.error("Models not loaded. Please load the models first.")
+    #     return None
+    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, "cuda", disable_mmap=False
+            )
+            clip_l = flux_utils.load_clip_l(CLIP_L_PATH, torch.bfloat16, "cuda", disable_mmap=False)
+            clip_l.eval()
+            t5xxl = flux_utils.load_t5xxl(T5XXL_PATH, torch.bfloat16, "cuda", disable_mmap=False)
+            t5xxl.eval()
+            ae = flux_utils.load_ae(AE_PATH, torch.bfloat16, "cuda", disable_mmap=False)
+
+        # Load LoRA weights
+        multiplier = 1.0
+        weights_sd = load_file(LORA_WEIGHTS_PATH)
+        lora_model, _ = lora_flux.create_network_from_weights(multiplier, None, ae, [clip_l, t5xxl], model, weights_sd, True)
+        lora_model.apply_to([clip_l, t5xxl], model)
+        info = lora_model.load_state_dict(weights_sd, strict=True)
+        logger.info(f"Loaded LoRA weights from {LORA_WEIGHTS_PATH}: {info}")
+        lora_model.eval()
+
+        logger.info("Models loaded successfully.")
+        # return "Models loaded successfully. Using Recraft: {}".format(selected_model)
+
+    except Exception as e:
+        logger.error(f"Error loading models: {e}")
+        return f"Error loading models: {e}"
 
     model_path = model_paths[recraft_model]
     frame_num = model_path['Frame']