Updates

app.py

@@ -27,7 +27,7 @@ base_model = "black-forest-labs/FLUX.1-dev"
 taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).to(device)
 good_vae = AutoencoderKL.from_pretrained(base_model, subfolder="vae", torch_dtype=dtype).to(device)
 pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=dtype, vae=taef1).to(device)
-ipe_i2i = FluxImg2ImgPipeline.from_pretrained(
+pipe_i2i = FluxImg2ImgPipeline.from_pretrained(
     base_model,
     vae=good_vae,
     transformer=pipe.transformer,
@@ -121,7 +121,7 @@ def generate_image_to_image(prompt_mash, image_input_path, image_strength, steps
     ).images[0]
     return final_image
 
-def run_lora(prompt, cfg_scale, steps, randomize_seed, seed, width, height, lora_scale, progress=gr.Progress(track_tqdm=True)):
+def run_lora(prompt, cfg_scale, steps, selected_index, randomize_seed, seed, width, height, lora_scale, progress=gr.Progress(track_tqdm=True)):
     global selected_lora_index
     if selected_lora_index is None:
         raise gr.Error("You must select a LoRA before proceeding.")
@@ -143,20 +143,13 @@ def run_lora(prompt, cfg_scale, steps, randomize_seed, seed, width, height, lora
 
     with calculateDuration("Unloading LoRA"):
         pipe.unload_lora_weights()
-        pipe_i2i.unload_lora_weights()
 
     # Load LoRA weights
     with calculateDuration(f"Loading LoRA weights for {selected_lora['title']}"):
-        if image_input is not None:
-            if "weights" in selected_lora:
-                pipe_i2i.load_lora_weights(lora_path, weight_name=selected_lora["weights"])
-            else:
-                pipe_i2i.load_lora_weights(lora_path)
+        if "weights" in selected_lora:
+            pipe.load_lora_weights(lora_path, weight_name=selected_lora["weights"])
         else:
-            if "weights" in selected_lora:
-                pipe.load_lora_weights(lora_path, weight_name=selected_lora["weights"])
-            else:
-                pipe.load_lora_weights(lora_path)
+            pipe.load_lora_weights(lora_path)
 
     # Set random seed for reproducibility
     with calculateDuration("Randomizing seed"):
@@ -168,16 +161,16 @@ def run_lora(prompt, cfg_scale, steps, randomize_seed, seed, width, height, lora
         yield final_image, seed, gr.update(visible=False)
     else:
         image_generator = generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scale, progress)
-        # Consume the generator to get the final image
-        final_image = None
-        step_counter = 0
-        for image in image_generator:
-            step_counter += 1
-            final_image = image
-            progress_bar = f'Generating image... Step {step_counter}/{steps}'
-            yield image, seed, gr.update(visible=True, value=progress_bar)
-        
-        yield final_image, seed, gr.update(visible=False)
+    # Consume the generator to get the final image
+    final_image = None
+    step_counter = 0
+    for image in image_generator:
+        step_counter += 1
+        final_image = image
+        progress_bar = f'Generating image... Step {step_counter}/{steps}'
+        yield image, seed, gr.update(visible=True, value=progress_bar)
+    
+    yield final_image, seed, gr.update(visible=False)
 
 # ...