main

gradio_app.py

@@ -12,61 +12,84 @@ from safetensors.torch import load_file
 from networks import lora_flux
 from library import flux_utils, flux_train_utils_recraft as flux_train_utils, strategy_flux
 import logging
+from huggingface_hub import login
+from huggingface_hub import hf_hub_download
+
 
 # Set up logger
 logger = logging.getLogger(__name__)
 logging.basicConfig(level=logging.DEBUG)
 
-# Ensure necessary devices are available
-print("torch.__version__: ", torch.__version__)
-device = "cuda" if torch.cuda.is_available() else "cpu"
-logger.info("device: {}".format(device))
-
 accelerator = Accelerator(mixed_precision='bf16', device_placement=True)
 
-# Model paths (replace these with your actual model paths)
-BASE_FLUX_CHECKPOINT="/tiamat-NAS/songyiren/FYP/liucheng/sd-scripts/MergeModel/6_Portrait/6_Portrait.safetensors"
-LORA_WEIGHTS_PATH="/tiamat-NAS/songyiren/FYP/liucheng/sd-scripts/RecraftModel/6_Portrait/6_Portrait-step00025000.safetensors"
-CLIP_L_PATH="/tiamat-NAS/hailong/storage_backup/models/stabilityai/stable-diffusion-3-medium/text_encoders/clip_l.safetensors"
-T5XXL_PATH="/tiamat-NAS/hailong/storage_backup/models/stabilityai/stable-diffusion-3-medium/text_encoders/t5xxl_fp16.safetensors"
-AE_PATH="/tiamat-vePFS/share_data/storage/huggingface/models/black-forest-labs/FLUX.1-dev/ae.safetensors"
-
-from huggingface_hub import login
 hf_token = os.getenv("HF_TOKEN")
 login(token=hf_token)
 
-from huggingface_hub import hf_hub_download
-repo_id = "Kijai/flux-fp8"
-file_name = "flux1-dev-fp8.safetensors"
-BASE_FLUX_CHECKPOINT = hf_hub_download(repo_id=repo_id, filename=file_name)
-print(f"Downloaded BASE_FLUX_CHECKPOINT saved at: {BASE_FLUX_CHECKPOINT}")
-
-from huggingface_hub import hf_hub_download
-repo_id = "comfyanonymous/flux_text_encoders"
-file_name = "clip_l.safetensors"
-CLIP_L_PATH = hf_hub_download(repo_id=repo_id, filename=file_name)
-print(f"Downloaded CLIP_L_PATH saved at: {CLIP_L_PATH}")
-
-from huggingface_hub import hf_hub_download
-repo_id = "comfyanonymous/flux_text_encoders"
-file_name = "t5xxl_fp8_e4m3fn.safetensors"
-T5XXL_PATH = hf_hub_download(repo_id=repo_id, filename=file_name)
-print(f"Downloaded T5XXL_PATH saved at: {T5XXL_PATH}")
-
-from huggingface_hub import hf_hub_download
-repo_id = "black-forest-labs/FLUX.1-dev"
-file_name = "ae.safetensors"
-AE_PATH = hf_hub_download(repo_id=repo_id, filename=file_name)
-print(f"Downloaded AE_PATH saved at: {AE_PATH}")
-
-from huggingface_hub import hf_hub_download
-repo_id = "showlab/makeanything"
-file_name = "recraft/recraft_9f_lego.safetensors"
-LORA_WEIGHTS_PATH = hf_hub_download(repo_id=repo_id, filename=file_name)
-print(f"Downloaded LORA_WEIGHTS_PATH saved at: {LORA_WEIGHTS_PATH}")
+# Model paths dynamically retrieved using selected model
+model_paths = {
+    'Wood Sculpture': {
+        'BASE_FLUX_CHECKPOINT': "Kijai/flux-fp8",
+        'BASE_FILE': "flux1-dev-fp8.safetensors",
+        'LORA_REPO': "showlab/makeanything",
+        'LORA_FILE': "recraft/recraft_4f_wood_sculpture.safetensors"
+    },
+    'LEGO': {
+        'BASE_FLUX_CHECKPOINT': "Kijai/flux-fp8",
+        'BASE_FILE': "flux1-dev-fp8.safetensors",
+        'LORA_REPO': "showlab/makeanything",
+        'LORA_FILE': "recraft/recraft_9f_lego.safetensors"
+    },
+    'Sketch': {
+        'BASE_FLUX_CHECKPOINT': "Kijai/flux-fp8",
+        'BASE_FILE': "flux1-dev-fp8.safetensors",
+        'LORA_REPO': "showlab/makeanything",
+        'LORA_FILE': "recraft/recraft_9f_sketch.safetensors"
+    },
+    'Portrait': {
+        'BASE_FLUX_CHECKPOINT': "Kijai/flux-fp8",
+        'BASE_FILE': "flux1-dev-fp8.safetensors",
+        'LORA_REPO': "showlab/makeanything",
+        'LORA_FILE': "recraft/recraft_9f_portrait.safetensors"
+    }
+}
+
+# Common paths
+clip_repo_id = "comfyanonymous/flux_text_encoders"
+t5xxl_file = "t5xxl_fp8_e4m3fn.safetensors"
+clip_l_file = "clip_l.safetensors"
+ae_repo_id = "black-forest-labs/FLUX.1-dev"
+ae_file = "ae.safetensors"
+
+
+# Model placeholders
+model = None
+clip_l = None
+t5xxl = None
+ae = None
+lora_model = None
+
+# Function to load a file from Hugging Face Hub
+def download_file(repo_id, file_name):
+    return hf_hub_download(repo_id=repo_id, filename=file_name)
 
 # Load model function
-def load_target_model():
+def load_target_model(selected_model):
+    global model, clip_l, t5xxl, ae, lora_model
+
+    # Fetch paths based on the selected model
+    model_path = model_paths[selected_model]
+    base_checkpoint_repo = model_path['BASE_FLUX_CHECKPOINT']
+    base_checkpoint_file = model_path['BASE_FILE']
+    lora_repo = model_path['LORA_REPO']
+    lora_file = model_path['LORA_FILE']
+
+    # Download necessary files
+    BASE_FLUX_CHECKPOINT = download_file(base_checkpoint_repo, base_checkpoint_file)
+    CLIP_L_PATH = download_file(clip_repo_id, clip_l_file)
+    T5XXL_PATH = download_file(clip_repo_id, t5xxl_file)
+    AE_PATH = download_file(ae_repo_id, ae_file)
+    LORA_WEIGHTS_PATH = download_file(lora_repo, lora_file)
+
     logger.info("Loading models...")
     try:
         _, model = flux_utils.load_flow_model(
@@ -108,32 +131,36 @@ class ResizeWithPadding:
 
 # The function to generate image from a prompt and conditional image
 @spaces.GPU(duration=180)
-def infer(prompt, sample_image, frame_num, seed=0, randomize_seed=False):
+def infer(prompt, sample_image, frame_num, 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(f"Started generating image with prompt: {prompt}")
-    
-    # Load models
-    model, [clip_l, t5xxl], ae = load_target_model()
 
+    lora_model.to("cuda")
+    
     model.eval()
     clip_l.eval()
     t5xxl.eval()
     ae.eval()
 
-    # LoRA
-    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()
-    lora_model.to(device)
-
-    # Process the seed
-    if randomize_seed:
-        seed = random.randint(0, np.iinfo(np.int32).max)
+    # # Load models
+    # model, [clip_l, t5xxl], ae = load_target_model()
+
+    # # LoRA
+    # 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()
+    # lora_model.to(device)
+
     logger.debug(f"Using seed: {seed}")
 
     # Preprocess the conditional image
@@ -199,9 +226,6 @@ def infer(prompt, sample_image, frame_num, seed=0, randomize_seed=False):
 
     model.to(device)
 
-    # import pdb
-    # pdb.set_trace()
-
     # Run the denoising process
     with accelerator.autocast(), torch.no_grad():
         x = flux_train_utils.denoise(
@@ -232,39 +256,39 @@ with gr.Blocks() as demo:
     gr.Markdown("## FLUX Image Generation")
     
     with gr.Row():
+        # Dropdown for selecting the recraft model
+        recraft_model = gr.Dropdown(
+            label="Select Recraft Model", 
+            choices=["Wood Sculpture", "LEGO", "Sketch", "Portrait"], 
+            value="Wood Sculpture"
+        )
+
         # Input for the prompt
         prompt = gr.Textbox(label="Prompt", placeholder="Enter your prompt here", lines=1)
-        
+
         # File upload for image
         sample_image = gr.Image(label="Upload a Conditional Image", type="pil")
-        
+
         # Frame number selection
-        frame_num = gr.Radio([4, 9], label="Select Frame Number", value=9)
-        
-        # Seed and randomize seed options
+        frame_num = gr.Radio([4, 9], label="Select Frame Number", value=4)
+
+        # Seed
         seed = gr.Slider(0, np.iinfo(np.int32).max, step=1, label="Seed", value=0)
-        randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
-        
+
+        # Load Model Button
+        load_button = gr.Button("Load Model")
+
         # Run Button
         run_button = gr.Button("Generate Image")
-        
+
         # Output result
         result_image = gr.Image(label="Generated Image")
 
-    run_button.click(
-        fn=infer,
-        inputs=[prompt, sample_image, frame_num, seed, randomize_seed],
-        outputs=[result_image]
-    )
+    # Load model button action
+    load_button.click(fn=load_target_model, inputs=[recraft_model], outputs=[])
+
+    # Run Button 
+    run_button.click(fn=infer, inputs=[prompt, sample_image, frame_num, seed], outputs=[result_image])
 
     # Launch the Gradio app
     demo.launch()
-
-
-# prompt = "1girl"
-# sample_image = Image.open("/tiamat-NAS/songyiren/FYP/liucheng/sd-scripts/MergeModel/test/1.png")  # 使用一个测试图像
-# frame_num = 9
-# seed = 42
-# randomize_seed = False
-# result = infer(prompt, sample_image, frame_num, seed, randomize_seed)
-# result.save('asy_results/generated_image.png')