main

gradio_app.py

@@ -94,35 +94,33 @@ def load_target_model(selected_model):
     AE_PATH = download_file(ae_repo_id, ae_file)
     LORA_WEIGHTS_PATH = download_file(lora_repo, lora_file)
 
-    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}"
+    logger.info("Loading models...")
+    try:
+        if model is None is None or clip_l is None or t5xxl is None or ae is None:
+            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 flux & LoRA weights
+        _, model = flux_utils.load_flow_model(
+            BASE_FLUX_CHECKPOINT, torch.float8_e4m3fn, "cpu", disable_mmap=False
+        )
+        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:
@@ -156,37 +154,10 @@ 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
-    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}"
+    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
 
     model_path = model_paths[recraft_model]
     frame_num = model_path['Frame']
@@ -317,7 +288,7 @@ def infer(prompt, sample_image, recraft_model, seed=0):
 
 # Gradio interface
 with gr.Blocks() as demo:
-    gr.Markdown("## FLUX Image Generation")
+    gr.Markdown("## Recraft Generation")
 
     with gr.Row():
         with gr.Column(scale=1):