Update app.py

app.py

@@ -113,45 +113,21 @@ taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).
 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)
 
-# 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
-    )
-    generated_image = pipe_i2i(prompt="A generated image")
-    
-    if isinstance(generated_image, Image.Image):
-        return generated_image  # Gradio can handle PIL Image objects
-    elif isinstance(generated_image, str):  # If it returns a file path, return that
-        return generated_image
-    else:
-        return None
+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
+)
     
 MAX_SEED = 2**32 - 1
 
@@ -652,22 +628,7 @@ 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,
@@ -702,7 +663,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, enable_enhancement_checkbox, enhancement_weight_slider, lora_scale_1, lora_scale_2, randomize_seed, seed, width, height, loras_state],
+        inputs=[prompt, input_image, image_strength, cfg_scale, steps, selected_indices, 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),