Update app.py

app.py

@@ -27,7 +27,7 @@ clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch16")
 longformer_tokenizer = LongformerTokenizer.from_pretrained("allenai/longformer-base-4096")
 longformer_model = LongformerModel.from_pretrained("allenai/longformer-base-4096")
 
-# Load prompts for randomization
+#Load prompts for randomization
 df = pd.read_csv('prompts.csv', header=None)
 prompt_values = df.values.flatten()
 
@@ -42,7 +42,13 @@ 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 = DiffusionPipeline.from_pretrained(base_model, torch_dtype=dtype, vae=taef1,
+    transformer=pipe.transformer,
+    text_encoder=pipe.text_encoder,
+    tokenizer=pipe.tokenizer,
+    text_encoder_2=pipe.text_encoder_2,
+    tokenizer_2=pipe.tokenizer_2,
+).to(device)
 
 MAX_SEED = 2**32 - 1
 
@@ -470,9 +476,10 @@ def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, progress)
         ):
             print("Image generated successfully.")  # Debugging statement
             yield img
+            return final_image
 
 @spaces.GPU(duration=75)
-def run_lora(prompt, cfg_scale, steps, selected_indices, lora_scale_1, lora_scale_2, randomize_seed, seed, width, height, loras_state, progress=gr.Progress(track_tqdm=True)):
+def run_lora(prompt, cfg_scale, steps, selected_indices, lora_scale_1, lora_scale_2, lora_scale_3, lora_scale_4, randomize_seed, seed, width, height, loras_state, 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
 
@@ -498,8 +505,7 @@ def run_lora(prompt, cfg_scale, steps, selected_indices, lora_scale_1, lora_scal
     # Unload previous LoRA weights
     with calculateDuration("Unloading LoRA"):
         pipe.unload_lora_weights()
-        pipe_i2i.unload_lora_weights()
-        
+##        pipe_i2i.unload_lora_weights()        
     print("Active adapters before loading new LoRAs:", pipe.get_active_adapters())
     
     # Load LoRA weights with respective scales
@@ -514,7 +520,7 @@ def run_lora(prompt, cfg_scale, steps, selected_indices, lora_scale_1, lora_scal
             lora_path = lora['repo']
             weight_name = lora.get("weights")
             print(f"Lora Path: {lora_path}")
-            pipe_to_use = pipe_i2i if image_input is not None else pipe
+##            pipe_to_use = pipe_i2i if image_input is not None else pipe
             pipe_to_use.load_lora_weights(
                 lora_path, 
                 weight_name=weight_name if weight_name else None,
@@ -523,9 +529,9 @@ def run_lora(prompt, cfg_scale, steps, selected_indices, lora_scale_1, lora_scal
             )
         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:
+##        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("Active adapters after loading new LoRAs:", pipe.get_active_adapters())
@@ -539,8 +545,8 @@ def run_lora(prompt, cfg_scale, steps, selected_indices, lora_scale_1, lora_scal
     # 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)
-            yield final_image, seed, gr.update(visible=False)
+            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=True)
         else:
             image_generator = generate_image(prompt_mash, steps, seed, cfg_scale, width, height, progress)
             # Consume the generator to get the final image