Update app.py

app.py

@@ -112,20 +112,43 @@ 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)
-pipe_i2i = AutoPipelineForImage2Image.from_pretrained(
-    base_model,
-    vae=good_vae,
-    transformer=pipe.transformer,
-    text_encoder=pipe.text_encoder,
-    tokenizer=pipe.tokenizer,
-    text_encoder_2=pipe.text_encoder_2,
-    tokenizer_2=pipe.tokenizer_2,
-    torch_dtype=dtype
-)
+
+# Function to dynamically merge models
+def merge_models(base_model, enhancement_model, alpha=0.7):
+    for base_param, enhance_param in zip(base_model.parameters(), enhancement_model.parameters()):
+        base_param.data = alpha * base_param.data + (1 - alpha) * enhance_param.data
+    return base_model
+
+# Gradio interface function
+def process_image(enable_enhancement, weight_slider):
+    # Load enhancement model if enabled
+    if enable_enhancement:
+        enhancement_model_path = "xey/sldr_flux_nsfw_v2-studio"
+        try:
+            enhancement_model = AutoencoderKL.from_pretrained(enhancement_model_path, torch_dtype=dtype).to(device)
+            # Merge with the base VAE using the weight from the slider
+            merged_vae = merge_models(good_vae, enhancement_model, alpha=weight_slider)
+        except Exception as e:
+            return f"Failed to load or merge enhancement model: {e}"
+    else:
+        merged_vae = good_vae  # Use the base VAE if no enhancement is enabled
+
+    # Create the image pipeline with the updated VAE
+    pipe_i2i = AutoPipelineForImage2Image.from_pretrained(
+        base_model,
+        vae=merged_vae,
+        transformer=pipe.transformer,
+        text_encoder=pipe.text_encoder,
+        tokenizer=pipe.tokenizer,
+        text_encoder_2=pipe.text_encoder_2,
+        tokenizer_2=pipe.tokenizer_2,
+        torch_dtype=dtype
+    )
 
 MAX_SEED = 2**32 - 1
 
 pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)
+return "Pipeline updated with the selected enhancement model and weight."
 
 class calculateDuration:
     def __init__(self, activity_name=""):
@@ -622,11 +645,28 @@ with gr.Blocks(css=css, delete_cache=(60, 60)) as app:
                 with gr.Row():
                     randomize_seed = gr.Checkbox(True, label="Randomize seed")
                     seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0, randomize=True)
+        # Gradio UI Elements for Enhancement Model and Weight Slider
+            with gr.Row():
+                with gr.Column():
+                    enable_enhancement_checkbox = gr.Checkbox(
+                    label="Enable Enhancement Model", 
+                    value=False, 
+                    elem_id="enable_enhancement_checkbox"
+            )
+                enhancement_weight_slider = gr.Slider(
+                label="Weight for Enhancement Model", 
+                minimum=0.0, 
+                maximum=1.0, 
+                step=0.05, 
+                value=0.75,  # Default weight
+                elem_id="enhancement_weight_slider"
+            )
 
     gallery.select(
         update_selection,
         inputs=[selected_indices, loras_state, width, height],
         outputs=[prompt, selected_info_1, selected_info_2, selected_indices, lora_scale_1, lora_scale_2, width, height, lora_image_1, lora_image_2])
+    
     remove_button_1.click(
         remove_lora_1,
         inputs=[selected_indices, loras_state],
@@ -655,7 +695,7 @@ with gr.Blocks(css=css, delete_cache=(60, 60)) as app:
     gr.on(
         triggers=[generate_button.click, prompt.submit],
         fn=run_lora,
-        inputs=[prompt, input_image, image_strength, cfg_scale, steps, selected_indices, lora_scale_1, lora_scale_2, randomize_seed, seed, width, height, loras_state],
+        inputs=[prompt, input_image, image_strength, cfg_scale, steps, selected_indices, enable_enhancement_checkbox, enhancement_weight_slider, lora_scale_1, lora_scale_2, randomize_seed, seed, width, height, loras_state],
         outputs=[result, seed, progress_bar]
     ).then(
         fn=lambda x, history: update_history(x, history),