First commit

flux_app/enhance.py

@@ -53,4 +53,3 @@ def generate(message, max_new_tokens=256, temperature=0.9, top_p=0.95, repetitio
         output += token_text
         yield output.strip('</s>')
     return output.strip('</s>')
-

flux_app/{enhance_v2.py → enhance_v1.py}

@@ -53,3 +53,4 @@ def generate(message, max_new_tokens=256, temperature=0.9, top_p=0.95, repetitio
         output += token_text
         yield output.strip('</s>')
     return output.strip('</s>')
+

flux_app/frontend.py

@@ -15,8 +15,9 @@ from flux_app.lora_handling import (
     unload_lora_weights, load_lora_weights_into_pipeline, update_selection
 )
 from flux_app.utilities import randomize_seed_if_needed, calculateDuration  # Absolute import
-# Import the prompt enhancer generate function from the new module
-from flux_app.enhance import generate
+
+# Import the prompt enhancer function
+from flux_app.enhance import generate as enhance_generate
 
 # Dummy loras data for initial UI setup.
 initial_loras = [
@@ -103,22 +104,29 @@ class Frontend:
             pass
 
     @spaces.GPU(duration=100)
-    def run_lora(self, prompt, image_input, image_strength, cfg_scale, steps, selected_index, randomize_seed, seed, width, height, lora_scale, use_enhancer, progress=gr.Progress(track_tqdm=True)):
-        # If prompt enhancer is enabled, generate the enhanced prompt.
+    def run_lora(self, prompt, image_input, image_strength, cfg_scale, steps, selected_index,
+                 randomize_seed, seed, width, height, lora_scale, use_enhancer,
+                 progress=gr.Progress(track_tqdm=True)):
+        seed = randomize_seed_if_needed(randomize_seed, seed, MAX_SEED)
+        # Prepare the initial prompt (using LoRA info if needed)
+        prompt_mash = prepare_prompt(prompt, selected_index, self.loras)
+        enhanced_text = ""
+        
+        # If prompt enhancer is enabled, first run it to improve the prompt.
         if use_enhancer:
-            enhanced_prompt = ""
-            # Generate the enhanced prompt (consume the generator to get the final result)
-            for chunk in generate(prompt):
-                enhanced_prompt = chunk
-            prompt_used = enhanced_prompt
+            # Stream the enhanced prompt (this will update the enhanced prompt textbox)
+            for enhanced_chunk in enhance_generate(prompt_mash):
+                enhanced_text = enhanced_chunk
+                # Yield an update with no image yet and the current enhanced prompt.
+                yield None, seed, gr.update(visible=False), enhanced_text
+            # Use the final enhanced prompt as the prompt for image generation.
+            prompt_mash = enhanced_text
         else:
-            enhanced_prompt = ""
-            prompt_used = prompt
-
-        seed = randomize_seed_if_needed(randomize_seed, seed, MAX_SEED)
-        prompt_mash = prepare_prompt(prompt_used, selected_index, self.loras)
+            # Ensure the enhanced prompt textbox remains cleared.
+            enhanced_text = ""
+        
+        # Continue with the image generation process.
         selected_lora = self.loras[selected_index]
-
         unload_lora_weights(self.model_manager.pipe, self.model_manager.pipe_i2i)
         pipe_to_use = self.model_manager.pipe_i2i if image_input is not None else self.model_manager.pipe
         load_lora_weights_into_pipeline(pipe_to_use, selected_lora["repo"], selected_lora.get("weights"))
@@ -127,20 +135,21 @@ class Frontend:
             final_image = self.model_manager.generate_image_to_image(
                 prompt_mash, image_input, image_strength, steps, cfg_scale, width, height, lora_scale, seed
             )
-            yield final_image, seed, gr.update(visible=False), enhanced_prompt
+            yield final_image, seed, gr.update(visible=False), enhanced_text
         else:
-            image_generator = self.model_manager.generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scale)
-            final_image = None
-            step_counter = 0
-            for image in image_generator:
+             image_generator = self.model_manager.generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scale)
+             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), enhanced_prompt
+                yield image, seed, gr.update(value=progress_bar, visible=True), enhanced_text
 
-            yield final_image, seed, gr.update(value=progress_bar, visible=False), enhanced_prompt
+             yield final_image, seed, gr.update(value=progress_bar, visible=False), enhanced_text
 
     def create_ui(self):
+        # Using a base theme for a clean and professional look.
         with gr.Blocks(theme=gr.themes.Base(), css=self.css, title="Flux LoRA Generation") as app:
             title = gr.HTML(
                 """<h1>Flux LoRA Generation</h1>""",
@@ -189,11 +198,11 @@ class Frontend:
                             randomize_seed = gr.Checkbox(True, label="Randomize seed")
                             seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0, randomize=True)
                             lora_scale = gr.Slider(label="LoRA Scale", minimum=0, maximum=3, step=0.01, value=0.95)
-                    # Prompt Enhancer Section
-                    with gr.Group():
-                        use_enhancer = gr.Checkbox(label="Use Prompt Enhancer", value=True)
-                        show_enhanced_prompt = gr.Checkbox(label="Display Enhanced Prompt", value=False)
-                        enhanced_prompt_box = gr.Textbox(label="Enhanced Prompt", lines=3, visible=False)
+                        with gr.Row():
+                            use_enhancer = gr.Checkbox(value=False, label="Use Prompt Enhancer")
+                            show_enhanced_prompt = gr.Checkbox(value=False, label="Display Enhanced Prompt")
+                    # Enhanced prompt textbox (hidden by default)
+                    enhanced_prompt_box = gr.Textbox(label="Enhanced Prompt", visible=False)
 
             gallery.select(
                 update_selection,
@@ -218,7 +227,8 @@ class Frontend:
             gr.on(
                 triggers=[generate_button.click, prompt.submit],
                 fn=self.run_lora,
-                inputs=[prompt, input_image, image_strength, cfg_scale, steps, selected_index, randomize_seed, seed, width, height, lora_scale, use_enhancer],
+                inputs=[prompt, input_image, image_strength, cfg_scale, steps, selected_index,
+                        randomize_seed, seed, width, height, lora_scale, use_enhancer],
                 outputs=[result, seed, progress_bar, enhanced_prompt_box]
             )
 
@@ -233,4 +243,4 @@ if __name__ == "__main__":
     frontend = Frontend(model_manager)
     app = frontend.create_ui()
     app.queue()
-    app.launch(debug=True)
+    app.launch()

flux_app/{frontend_v2.py → frontend_nw.py}

@@ -2,7 +2,7 @@
 import gradio as gr
 import sys
 import os
-
+import spaces
 # Add the parent directory to sys.path
 parent_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
 sys.path.insert(0, parent_dir)
@@ -15,9 +15,8 @@ from flux_app.lora_handling import (
     unload_lora_weights, load_lora_weights_into_pipeline, update_selection
 )
 from flux_app.utilities import randomize_seed_if_needed, calculateDuration  # Absolute import
-import spaces
-# Import the prompt enhancer function
-from flux_app.enhance import generate as enhance_generate
+# Import the prompt enhancer generate function from the new module
+from flux_app.enhance import generate
 
 # Dummy loras data for initial UI setup.
 initial_loras = [
@@ -104,29 +103,22 @@ class Frontend:
             pass
 
     @spaces.GPU(duration=100)
-    def run_lora(self, prompt, image_input, image_strength, cfg_scale, steps, selected_index,
-                 randomize_seed, seed, width, height, lora_scale, use_enhancer,
-                 progress=gr.Progress(track_tqdm=True)):
-        seed = randomize_seed_if_needed(randomize_seed, seed, MAX_SEED)
-        # Prepare the initial prompt (using LoRA info if needed)
-        prompt_mash = prepare_prompt(prompt, selected_index, self.loras)
-        enhanced_text = ""
-        
-        # If prompt enhancer is enabled, first run it to improve the prompt.
+    def run_lora(self, prompt, image_input, image_strength, cfg_scale, steps, selected_index, randomize_seed, seed, width, height, lora_scale, use_enhancer, progress=gr.Progress(track_tqdm=True)):
+        # If prompt enhancer is enabled, generate the enhanced prompt.
         if use_enhancer:
-            # Stream the enhanced prompt (this will update the enhanced prompt textbox)
-            for enhanced_chunk in enhance_generate(prompt_mash):
-                enhanced_text = enhanced_chunk
-                # Yield an update with no image yet and the current enhanced prompt.
-                yield None, seed, gr.update(visible=False), enhanced_text
-            # Use the final enhanced prompt as the prompt for image generation.
-            prompt_mash = enhanced_text
+            enhanced_prompt = ""
+            # Generate the enhanced prompt (consume the generator to get the final result)
+            for chunk in generate(prompt):
+                enhanced_prompt = chunk
+            prompt_used = enhanced_prompt
         else:
-            # Ensure the enhanced prompt textbox remains cleared.
-            enhanced_text = ""
-        
-        # Continue with the image generation process.
+            enhanced_prompt = ""
+            prompt_used = prompt
+
+        seed = randomize_seed_if_needed(randomize_seed, seed, MAX_SEED)
+        prompt_mash = prepare_prompt(prompt_used, selected_index, self.loras)
         selected_lora = self.loras[selected_index]
+
         unload_lora_weights(self.model_manager.pipe, self.model_manager.pipe_i2i)
         pipe_to_use = self.model_manager.pipe_i2i if image_input is not None else self.model_manager.pipe
         load_lora_weights_into_pipeline(pipe_to_use, selected_lora["repo"], selected_lora.get("weights"))
@@ -135,21 +127,20 @@ class Frontend:
             final_image = self.model_manager.generate_image_to_image(
                 prompt_mash, image_input, image_strength, steps, cfg_scale, width, height, lora_scale, seed
             )
-            yield final_image, seed, gr.update(visible=False), enhanced_text
+            yield final_image, seed, gr.update(visible=False), enhanced_prompt
         else:
-             image_generator = self.model_manager.generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scale)
-             final_image = None
-             step_counter = 0
-             for image in image_generator:
+            image_generator = self.model_manager.generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scale)
+            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), enhanced_text
+                yield image, seed, gr.update(value=progress_bar, visible=True), enhanced_prompt
 
-             yield final_image, seed, gr.update(value=progress_bar, visible=False), enhanced_text
+            yield final_image, seed, gr.update(value=progress_bar, visible=False), enhanced_prompt
 
     def create_ui(self):
-        # Using a base theme for a clean and professional look.
         with gr.Blocks(theme=gr.themes.Base(), css=self.css, title="Flux LoRA Generation") as app:
             title = gr.HTML(
                 """<h1>Flux LoRA Generation</h1>""",
@@ -198,11 +189,11 @@ class Frontend:
                             randomize_seed = gr.Checkbox(True, label="Randomize seed")
                             seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0, randomize=True)
                             lora_scale = gr.Slider(label="LoRA Scale", minimum=0, maximum=3, step=0.01, value=0.95)
-                        with gr.Row():
-                            use_enhancer = gr.Checkbox(value=False, label="Use Prompt Enhancer")
-                            show_enhanced_prompt = gr.Checkbox(value=False, label="Display Enhanced Prompt")
-                    # Enhanced prompt textbox (hidden by default)
-                    enhanced_prompt_box = gr.Textbox(label="Enhanced Prompt", visible=False)
+                    # Prompt Enhancer Section
+                    with gr.Group():
+                        use_enhancer = gr.Checkbox(label="Use Prompt Enhancer", value=True)
+                        show_enhanced_prompt = gr.Checkbox(label="Display Enhanced Prompt", value=False)
+                        enhanced_prompt_box = gr.Textbox(label="Enhanced Prompt", lines=3, visible=False)
 
             gallery.select(
                 update_selection,
@@ -227,8 +218,7 @@ class Frontend:
             gr.on(
                 triggers=[generate_button.click, prompt.submit],
                 fn=self.run_lora,
-                inputs=[prompt, input_image, image_strength, cfg_scale, steps, selected_index,
-                        randomize_seed, seed, width, height, lora_scale, use_enhancer],
+                inputs=[prompt, input_image, image_strength, cfg_scale, steps, selected_index, randomize_seed, seed, width, height, lora_scale, use_enhancer],
                 outputs=[result, seed, progress_bar, enhanced_prompt_box]
             )
 
@@ -243,4 +233,4 @@ if __name__ == "__main__":
     frontend = Frontend(model_manager)
     app = frontend.create_ui()
     app.queue()
-    app.launch()
+    app.launch(debug=True)