Update app.py

app.py

@@ -440,12 +440,7 @@ def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, progress)
     print("Generating image...")
     pipe.to("cuda")
     generator = torch.Generator(device="cuda").manual_seed(seed)
-    final_image = None  # Initialize final_image
     with calculateDuration("Generating image"):
-        # Ensure width and height are integers
-        width = int(width) if isinstance(width, (float, int)) else 1024
-        height = int(height) if isinstance(height, (float, int)) else 1024
-
         # Generate image
         for img in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
             prompt=prompt_mash,
@@ -458,17 +453,9 @@ def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, progress)
             output_type="pil",
             good_vae=good_vae,
         ):
-            print(f"Debug: Yielding image of type {type(img)}")  # Check image type
-            if isinstance(img, float):
-                print("Error: A float was returned instead of an image.")  # Log if img is a float
-                raise ValueError("Expected an image, but got a float.")  # Raise error if a float is found
             yield img
-            final_image = img  # Update final_image with the current image
-        return final_image
-
 
 def generate_image_to_image(prompt_mash, image_input_path, image_strength, steps, cfg_scale, width, height, seed):
-    print("Generating image from input...")
     pipe_i2i.to("cuda")
     generator = torch.Generator(device="cuda").manual_seed(seed)
     image_input = load_image(image_input_path)
@@ -484,22 +471,14 @@ def generate_image_to_image(prompt_mash, image_input_path, image_strength, steps
         joint_attention_kwargs={"scale": 1.0},
         output_type="pil",
     ).images[0]
-    if isinstance(final_image, float):
-        print("Error: Expected an image but got a float.")
-        raise ValueError("Expected an image, but got a float.")
     return final_image
 
-
 @spaces.GPU(duration=75)
-def run_lora(prompt, cfg_scale, steps, selected_info_1, selected_info_2, selected_info_3, selected_info_4, selected_indices, lora_scale_1, lora_scale_2, lora_scale_3, lora_scale_4, randomize_seed, seed, width, height, loras_state, image_input=None, progress=gr.Progress(track_tqdm=True)):
-    print("run_lora function called.")  # Debugging statement
-    print(f"Inputs received - Prompt: {prompt}, CFG Scale: {cfg_scale}, Steps: {steps}, Seed: {seed}, Width: {width}, Height: {height}")  # Debugging statement
-
+def run_lora(prompt, image_input, image_strength, cfg_scale, steps, selected_indices, lora_scale_1, lora_scale_2, randomize_seed, seed, width, height, loras_state, progress=gr.Progress(track_tqdm=True)):
     if not selected_indices:
         raise gr.Error("You must select at least one LoRA before proceeding.")
 
     selected_loras = [loras_state[idx] for idx in selected_indices]
-    print(f"Selected LoRAs: {selected_loras}")  # Debugging statement
 
     # Build the prompt with trigger words
     prepends = []
@@ -507,39 +486,27 @@ def run_lora(prompt, cfg_scale, steps, selected_info_1, selected_info_2, selecte
     for lora in selected_loras:
         trigger_word = lora.get('trigger_word', '')
         if trigger_word:
-            if lora.get("trigger_position") == "append":
+            if lora.get("trigger_position") == "prepend":
                 prepends.append(trigger_word)
             else:
                 appends.append(trigger_word)
     prompt_mash = " ".join(prepends + [prompt] + appends)
-    print("Prompt Mash: ", prompt_mash)  # Debugging statement
-
-    # Ensure valid width and height values
-    if width is None or isinstance(width, gr.Progress):  # Check for Gradio Progress object
-        width = 1024  # Default value
-    if height is None:
-        height = 1024  # Default value
-
-    # Set seed value
-    if seed is None or randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-
+    print("Prompt Mash: ", prompt_mash)
     # Unload previous LoRA weights
     with calculateDuration("Unloading LoRA"):
         pipe.unload_lora_weights()
-    print("Active adapters before loading new LoRAs:", pipe.get_active_adapters())
-    
+        pipe_i2i.unload_lora_weights()
+        
+    print(pipe.get_active_adapters())
     # Load LoRA weights with respective scales
     lora_names = []
     lora_weights = []
-    lora_scales = [lora_scale_1, lora_scale_2, lora_scale_3, lora_scale_4]  # List of scales
     with calculateDuration("Loading LoRA weights"):
         for idx, lora in enumerate(selected_loras):
             lora_name = f"lora_{idx}"
             lora_names.append(lora_name)
             print(f"Lora Name: {lora_name}")
-            # Use the scale based on the index, ensuring it does not exceed the list length
-            lora_weights.append(lora_scales[idx] if idx < len(lora_scales) else lora_scale_4)  # Default to last scale if out of bounds
+            lora_weights.append(lora_scale_1 if idx == 0 else lora_scale_2)
             lora_path = lora['repo']
             weight_name = lora.get("weights")
             print(f"Lora Path: {lora_path}")
@@ -550,47 +517,38 @@ def run_lora(prompt, cfg_scale, steps, selected_info_1, selected_info_2, selecte
                 low_cpu_mem_usage=True,
                 adapter_name=lora_name
             )
-    print("Loaded LoRAs:", lora_names)
-    print("Adapter weights:", lora_weights)
-    
+#            if image_input is not None: pipe_i2i = pipe_to_use
+#            else: pipe = pipe_to_use
+        print("Loaded LoRAs:", lora_names)
+        print("Adapter weights:", lora_weights)
+        if image_input is not None:
+            pipe_i2i.set_adapters(lora_names, adapter_weights=lora_weights)
+        else:
+            pipe.set_adapters(lora_names, adapter_weights=lora_weights)
+    print(pipe.get_active_adapters())
     # Set random seed for reproducibility
     with calculateDuration("Randomizing seed"):
-        print("Image Seed:", seed)
+        if randomize_seed:
+            seed = random.randint(0, MAX_SEED)
 
     # Generate image
-    try:
-        if image_input is not None:
-            final_image = generate_image_to_image(prompt_mash, image_input, image_strength, steps, cfg_scale, width, height, seed)
-            print(f"Debug: generate_image_to_image output type: {type(final_image)}")  # Debugging
-            if isinstance(final_image, float):
-                print("Error: Expected an image but got a float.")  # Log if it's a float
-                raise ValueError("Expected an image, but got a float.")  # Raise error if a float is found
-            yield final_image, seed, gr.update(visible=False)
-        else:
-            image_generator = generate_image(prompt_mash, steps, seed, cfg_scale, width, height, progress)
-            final_image = None
-            step_counter = 0
-            for image in image_generator:
-                print(f"Debug: generate_image yielded type: {type(image)}")  # Debugging
-                print(f"Debug: generate_image yielded value: {image}")  # Debugging
-                step_counter += 1
-                final_image = image
-                progress_bar = f'<div class="progress-container"><div class="progress-bar" style="--current: {step_counter}; --total: {steps};"></div></div>'
-                yield image, seed, gr.update(value=progress_bar, visible=True)
-
-            if final_image is None:
-                print("No final image generated.")  # Debugging statement
-            else:
-                print(f"Debug: final_image type: {type(final_image)}")  # Debugging
-                yield final_image, seed, gr.update(value=progress_bar, visible=False)
-
-    except Exception as e:
-        print(f"Error during image generation: {e}")  # Error handling
-        raise gr.Error("An error occurred during image generation.")
+    if image_input is not None:
+        final_image = generate_image_to_image(prompt_mash, image_input, image_strength, steps, cfg_scale, width, height, seed)
+        yield final_image, seed, gr.update(visible=False)
+    else:
+        image_generator = generate_image(prompt_mash, steps, seed, cfg_scale, width, height, 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'<div class="progress-container"><div class="progress-bar" style="--current: {step_counter}; --total: {steps};"></div></div>'
+            yield image, seed, gr.update(value=progress_bar, visible=True)
+        yield final_image, seed, gr.update(value=progress_bar, visible=False)
 
 run_lora.zerogpu = True
 
-
 def get_huggingface_safetensors(link):
     split_link = link.split("/")
     if len(split_link) == 4: