main

gradio_app.py

@@ -31,25 +31,29 @@ model_paths = {
         'BASE_FLUX_CHECKPOINT': "showlab/makeanything",
         'BASE_FILE': "flux_merge_lora/flux_merge_4f_wood-fp8_e4m3fn.safetensors",
         'LORA_REPO': "showlab/makeanything",
-        'LORA_FILE': "recraft/recraft_4f_wood_sculpture.safetensors"
+        'LORA_FILE': "recraft/recraft_4f_wood_sculpture.safetensors",
+        "Frame": 4
     },
     'LEGO': {
         'BASE_FLUX_CHECKPOINT': "showlab/makeanything",
         'BASE_FILE': "flux_merge_lora/flux_merge_9f_lego-fp8_e4m3fn.safetensors",
         'LORA_REPO': "showlab/makeanything",
-        'LORA_FILE': "recraft/recraft_9f_lego.safetensors"
+        'LORA_FILE': "recraft/recraft_9f_lego.safetensors",
+        "Frame": 9
     },
     'Sketch': {
         'BASE_FLUX_CHECKPOINT': "showlab/makeanything",
         'BASE_FILE': "flux_merge_lora/flux_merge_9f_portrait-fp8_e4m3fn.safetensors",
         'LORA_REPO': "showlab/makeanything",
-        'LORA_FILE': "recraft/recraft_9f_sketch.safetensors"
+        'LORA_FILE': "recraft/recraft_9f_sketch.safetensors",
+        "Frame": 9
     },
     'Portrait': {
         'BASE_FLUX_CHECKPOINT': "showlab/makeanything",
         'BASE_FILE': "flux_merge_lora/flux_merge_9f_sketch-fp8_e4m3fn.safetensors",
         'LORA_REPO': "showlab/makeanything",
-        'LORA_FILE': "recraft/recraft_9f_portrait.safetensors"
+        'LORA_FILE': "recraft/recraft_9f_portrait.safetensors",
+        "Frame": 9
     }
 }
 
@@ -92,14 +96,15 @@ def load_target_model(selected_model):
 
     logger.info("Loading models...")
     try:
-        _, 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)
+        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
@@ -148,12 +153,15 @@ 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):
+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
 
+    model_path = model_paths[selected_model]
+    frame_num = model_path['Frame']
+
     logger.info(f"Started generating image with prompt: {prompt}")
 
     lora_model.to("cuda")
@@ -288,9 +296,6 @@ with gr.Blocks() as demo:
         # 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
         seed = gr.Slider(0, np.iinfo(np.int32).max, step=1, label="Seed", value=0)
 
@@ -310,7 +315,7 @@ with gr.Blocks() as demo:
     load_button.click(fn=load_target_model, inputs=[recraft_model], outputs=[status_box])
 
     # Run Button 
-    run_button.click(fn=infer, inputs=[prompt, sample_image, frame_num, seed], outputs=[result_image])
+    run_button.click(fn=infer, inputs=[prompt, sample_image, recraft_model, seed], outputs=[result_image])
 
     # Launch the Gradio app
     demo.launch()