ruslanmv commited on
Commit
8236b2b
·
1 Parent(s): bf7e3bb
.gitignore ADDED
@@ -0,0 +1 @@
 
 
1
+ /backup
Flux-LoRA-Generation-Advanced.zip ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:2537b224b8b98c72939afaf580fd85d9c375d1f6e1f94b1c2f630f22fc0f03ce
3
+ size 26901
flux_app/backend.py CHANGED
@@ -7,8 +7,8 @@ from diffusers import (
7
  AutoPipelineForImage2Image,
8
  )
9
  from flux_app.config import DTYPE, DEVICE, BASE_MODEL, TAEF1_MODEL, MAX_SEED # Absolute import
10
- from flux_app.utilities import calculate_shift, retrieve_timesteps, load_image_from_path, calculateDuration # Absolute import
11
- from flux_app.lora_handling import flux_pipe_call_that_returns_an_iterable_of_images # Absolute import
12
  import time
13
  from huggingface_hub import login
14
 
@@ -21,17 +21,14 @@ class ModelManager:
21
 
22
  if hf_token:
23
  login(token=hf_token) # Log in with the provided token
24
- #else: # Optional: You could add a fallback to interactive login
25
- # login()
26
 
27
  self.initialize_models()
28
 
29
-
30
  def initialize_models(self):
31
  """Initializes the diffusion pipelines and autoencoders."""
32
- self.taef1 = AutoencoderTiny.from_pretrained(TAEF1_MODEL, torch_dtype=DTYPE, token=True).to(DEVICE)
33
- self.good_vae = AutoencoderKL.from_pretrained(BASE_MODEL, subfolder="vae", torch_dtype=DTYPE, token=True).to(DEVICE)
34
- self.pipe = DiffusionPipeline.from_pretrained(BASE_MODEL, torch_dtype=DTYPE, vae=self.taef1, token=True).to(DEVICE)
35
  self.pipe_i2i = AutoPipelineForImage2Image.from_pretrained(
36
  BASE_MODEL,
37
  vae=self.good_vae,
@@ -40,18 +37,17 @@ class ModelManager:
40
  tokenizer=self.pipe.tokenizer,
41
  text_encoder_2=self.pipe.text_encoder_2,
42
  tokenizer_2=self.pipe.tokenizer_2,
43
- torch_dtype=DTYPE,
44
- token=True
45
  )
46
- self.pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(self.pipe)
47
 
 
 
48
 
49
  def generate_image(self, prompt_mash, steps, seed, cfg_scale, width, height, lora_scale):
50
  """Generates an image using the text-to-image pipeline."""
51
  self.pipe.to(DEVICE)
52
  generator = torch.Generator(device=DEVICE).manual_seed(seed)
53
  with calculateDuration("Generating image"):
54
-
55
  for img in self.pipe.flux_pipe_call_that_returns_an_iterable_of_images(
56
  prompt=prompt_mash,
57
  num_inference_steps=steps,
@@ -83,4 +79,4 @@ class ModelManager:
83
  joint_attention_kwargs={"scale": lora_scale},
84
  output_type="pil",
85
  ).images[0]
86
- return final_image
 
7
  AutoPipelineForImage2Image,
8
  )
9
  from flux_app.config import DTYPE, DEVICE, BASE_MODEL, TAEF1_MODEL, MAX_SEED # Absolute import
10
+ from flux_app.utilities import calculate_shift, retrieve_timesteps, load_image_from_path, calculateDuration # Absolute import
11
+ from flux_app.lora_handling import flux_pipe_call_that_returns_an_iterable_of_images # Absolute import
12
  import time
13
  from huggingface_hub import login
14
 
 
21
 
22
  if hf_token:
23
  login(token=hf_token) # Log in with the provided token
 
 
24
 
25
  self.initialize_models()
26
 
 
27
  def initialize_models(self):
28
  """Initializes the diffusion pipelines and autoencoders."""
29
+ self.taef1 = AutoencoderTiny.from_pretrained(TAEF1_MODEL, torch_dtype=DTYPE).to(DEVICE)
30
+ self.good_vae = AutoencoderKL.from_pretrained(BASE_MODEL, subfolder="vae", torch_dtype=DTYPE).to(DEVICE)
31
+ self.pipe = DiffusionPipeline.from_pretrained(BASE_MODEL, torch_dtype=DTYPE, vae=self.taef1).to(DEVICE)
32
  self.pipe_i2i = AutoPipelineForImage2Image.from_pretrained(
33
  BASE_MODEL,
34
  vae=self.good_vae,
 
37
  tokenizer=self.pipe.tokenizer,
38
  text_encoder_2=self.pipe.text_encoder_2,
39
  tokenizer_2=self.pipe.tokenizer_2,
40
+ torch_dtype=DTYPE
 
41
  )
 
42
 
43
+ setattr(self.pipe, "flux_pipe_call_that_returns_an_iterable_of_images",
44
+ lambda *args, **kwargs: flux_pipe_call_that_returns_an_iterable_of_images(self.pipe, *args, **kwargs))
45
 
46
  def generate_image(self, prompt_mash, steps, seed, cfg_scale, width, height, lora_scale):
47
  """Generates an image using the text-to-image pipeline."""
48
  self.pipe.to(DEVICE)
49
  generator = torch.Generator(device=DEVICE).manual_seed(seed)
50
  with calculateDuration("Generating image"):
 
51
  for img in self.pipe.flux_pipe_call_that_returns_an_iterable_of_images(
52
  prompt=prompt_mash,
53
  num_inference_steps=steps,
 
79
  joint_attention_kwargs={"scale": lora_scale},
80
  output_type="pil",
81
  ).images[0]
82
+ return final_image
flux_app/enhance_v1.py DELETED
@@ -1,56 +0,0 @@
1
- # flux_app/enhance.py
2
- import time
3
- from huggingface_hub import InferenceClient
4
- import gradio as gr
5
-
6
- # Initialize the inference client with the new LLM
7
- client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")
8
-
9
- # Define the system prompt for enhancing user prompts
10
- SYSTEM_PROMPT = (
11
- "You are a prompt enhancer and your work is to enhance the given prompt under 100 words "
12
- "without changing the essence, only write the enhanced prompt and nothing else."
13
- )
14
-
15
- def format_prompt(message):
16
- """
17
- Format the input message using the system prompt and a timestamp to ensure uniqueness.
18
- """
19
- timestamp = time.time()
20
- formatted = (
21
- f"<s>[INST] SYSTEM: {SYSTEM_PROMPT} [/INST]"
22
- f"[INST] {message} {timestamp} [/INST]"
23
- )
24
- return formatted
25
-
26
- def generate(message, max_new_tokens=256, temperature=0.9, top_p=0.95, repetition_penalty=1.0):
27
- """
28
- Generate an enhanced prompt using the new LLM.
29
- This function yields intermediate results as they are generated.
30
- """
31
- temperature = float(temperature)
32
- if temperature < 1e-2:
33
- temperature = 1e-2
34
- top_p = float(top_p)
35
- generate_kwargs = {
36
- "temperature": temperature,
37
- "max_new_tokens": int(max_new_tokens),
38
- "top_p": top_p,
39
- "repetition_penalty": float(repetition_penalty),
40
- "do_sample": True,
41
- }
42
- formatted_prompt = format_prompt(message)
43
- stream = client.text_generation(
44
- formatted_prompt,
45
- **generate_kwargs,
46
- stream=True,
47
- details=True,
48
- return_full_text=False,
49
- )
50
- output = ""
51
- for response in stream:
52
- token_text = response.token.text
53
- output += token_text
54
- yield output.strip('</s>')
55
- return output.strip('</s>')
56
-
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
flux_app/enhance_v2.py DELETED
@@ -1,55 +0,0 @@
1
- # flux_app/enhance.py
2
- import time
3
- from huggingface_hub import InferenceClient
4
- import gradio as gr
5
-
6
- # Initialize the inference client with the new LLM
7
- client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")
8
-
9
- # Define the system prompt for enhancing user prompts
10
- SYSTEM_PROMPT = (
11
- "You are a prompt enhancer and your work is to enhance the given prompt under 100 words "
12
- "without changing the essence, only write the enhanced prompt and nothing else."
13
- )
14
-
15
- def format_prompt(message):
16
- """
17
- Format the input message using the system prompt and a timestamp to ensure uniqueness.
18
- """
19
- timestamp = time.time()
20
- formatted = (
21
- f"<s>[INST] SYSTEM: {SYSTEM_PROMPT} [/INST]"
22
- f"[INST] {message} {timestamp} [/INST]"
23
- )
24
- return formatted
25
-
26
- def generate(message, max_new_tokens=256, temperature=0.9, top_p=0.95, repetition_penalty=1.0):
27
- """
28
- Generate an enhanced prompt using the new LLM.
29
- This function yields intermediate results as they are generated.
30
- """
31
- temperature = float(temperature)
32
- if temperature < 1e-2:
33
- temperature = 1e-2
34
- top_p = float(top_p)
35
- generate_kwargs = {
36
- "temperature": temperature,
37
- "max_new_tokens": int(max_new_tokens),
38
- "top_p": top_p,
39
- "repetition_penalty": float(repetition_penalty),
40
- "do_sample": True,
41
- }
42
- formatted_prompt = format_prompt(message)
43
- stream = client.text_generation(
44
- formatted_prompt,
45
- **generate_kwargs,
46
- stream=True,
47
- details=True,
48
- return_full_text=False,
49
- )
50
- output = ""
51
- for response in stream:
52
- token_text = response.token.text
53
- output += token_text
54
- yield output.strip('</s>')
55
- return output.strip('</s>')
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
flux_app/frontend.py CHANGED
@@ -103,7 +103,7 @@ class Frontend:
103
  print("Warning: lora.py not found, using placeholder LoRAs.")
104
  pass
105
 
106
- @spaces.GPU(duration=100)
107
  def run_lora(self, prompt, image_input, image_strength, cfg_scale, steps, selected_index,
108
  randomize_seed, seed, width, height, lora_scale, use_enhancer,
109
  progress=gr.Progress(track_tqdm=True)):
 
103
  print("Warning: lora.py not found, using placeholder LoRAs.")
104
  pass
105
 
106
+ @spaces.GPU(duration=300)
107
  def run_lora(self, prompt, image_input, image_strength, cfg_scale, steps, selected_index,
108
  randomize_seed, seed, width, height, lora_scale, use_enhancer,
109
  progress=gr.Progress(track_tqdm=True)):
flux_app/frontend_nw.py DELETED
@@ -1,236 +0,0 @@
1
- # frontend.py
2
- import gradio as gr
3
- import sys
4
- import os
5
- import spaces
6
- # Add the parent directory to sys.path
7
- parent_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
8
- sys.path.insert(0, parent_dir)
9
- #print(sys.path) #DEBUG
10
-
11
- from flux_app.backend import ModelManager # Absolute import
12
- from flux_app.config import MAX_SEED # Absolute import
13
- from flux_app.lora_handling import (
14
- add_custom_lora, remove_custom_lora, prepare_prompt,
15
- unload_lora_weights, load_lora_weights_into_pipeline, update_selection
16
- )
17
- from flux_app.utilities import randomize_seed_if_needed, calculateDuration # Absolute import
18
- # Import the prompt enhancer generate function from the new module
19
- from flux_app.enhance import generate
20
-
21
- # Dummy loras data for initial UI setup.
22
- initial_loras = [
23
- {"image": "placeholder.jpg", "title": "Placeholder LoRA", "repo": "placeholder/repo", "weights": None, "trigger_word": ""},
24
- ]
25
-
26
- class Frontend:
27
- def __init__(self, model_manager: ModelManager):
28
- self.model_manager = model_manager
29
- self.loras = initial_loras
30
- self.load_initial_loras()
31
- self.css = self.define_css()
32
-
33
- def define_css(self):
34
- # A cleaner, professional CSS styling.
35
- return '''
36
- /* Title Styling */
37
- #title {
38
- text-align: center;
39
- margin-bottom: 20px;
40
- }
41
- #title h1 {
42
- font-size: 2.5rem;
43
- margin: 0;
44
- color: #333;
45
- }
46
- /* Button and Column Styling */
47
- #gen_btn {
48
- width: 100%;
49
- padding: 12px;
50
- font-weight: bold;
51
- border-radius: 5px;
52
- }
53
- #gen_column {
54
- display: flex;
55
- align-items: center;
56
- justify-content: center;
57
- }
58
- /* Gallery and List Styling */
59
- #gallery .grid-wrap {
60
- margin-top: 15px;
61
- }
62
- #lora_list {
63
- background-color: #f5f5f5;
64
- padding: 10px;
65
- border-radius: 4px;
66
- font-size: 0.9rem;
67
- }
68
- .card_internal {
69
- display: flex;
70
- align-items: center;
71
- height: 100px;
72
- margin-top: 10px;
73
- }
74
- .card_internal img {
75
- margin-right: 10px;
76
- }
77
- .styler {
78
- --form-gap-width: 0px !important;
79
- }
80
- /* Progress Bar Styling */
81
- .progress-container {
82
- width: 100%;
83
- height: 20px;
84
- background-color: #e0e0e0;
85
- border-radius: 10px;
86
- overflow: hidden;
87
- margin-bottom: 20px;
88
- }
89
- .progress-bar {
90
- height: 100%;
91
- background-color: #4f46e5;
92
- transition: width 0.3s ease-in-out;
93
- width: calc(var(--current) / var(--total) * 100%);
94
- }
95
- '''
96
-
97
- def load_initial_loras(self):
98
- try:
99
- from flux_app.lora import loras as loras_list # Absolute import
100
- self.loras = loras_list
101
- except ImportError:
102
- print("Warning: lora.py not found, using placeholder LoRAs.")
103
- pass
104
-
105
- @spaces.GPU(duration=100)
106
- 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)):
107
- # If prompt enhancer is enabled, generate the enhanced prompt.
108
- if use_enhancer:
109
- enhanced_prompt = ""
110
- # Generate the enhanced prompt (consume the generator to get the final result)
111
- for chunk in generate(prompt):
112
- enhanced_prompt = chunk
113
- prompt_used = enhanced_prompt
114
- else:
115
- enhanced_prompt = ""
116
- prompt_used = prompt
117
-
118
- seed = randomize_seed_if_needed(randomize_seed, seed, MAX_SEED)
119
- prompt_mash = prepare_prompt(prompt_used, selected_index, self.loras)
120
- selected_lora = self.loras[selected_index]
121
-
122
- unload_lora_weights(self.model_manager.pipe, self.model_manager.pipe_i2i)
123
- pipe_to_use = self.model_manager.pipe_i2i if image_input is not None else self.model_manager.pipe
124
- load_lora_weights_into_pipeline(pipe_to_use, selected_lora["repo"], selected_lora.get("weights"))
125
-
126
- if image_input is not None:
127
- final_image = self.model_manager.generate_image_to_image(
128
- prompt_mash, image_input, image_strength, steps, cfg_scale, width, height, lora_scale, seed
129
- )
130
- yield final_image, seed, gr.update(visible=False), enhanced_prompt
131
- else:
132
- image_generator = self.model_manager.generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scale)
133
- final_image = None
134
- step_counter = 0
135
- for image in image_generator:
136
- step_counter += 1
137
- final_image = image
138
- progress_bar = f'<div class="progress-container"><div class="progress-bar" style="--current: {step_counter}; --total: {steps};"></div></div>'
139
- yield image, seed, gr.update(value=progress_bar, visible=True), enhanced_prompt
140
-
141
- yield final_image, seed, gr.update(value=progress_bar, visible=False), enhanced_prompt
142
-
143
- def create_ui(self):
144
- with gr.Blocks(theme=gr.themes.Base(), css=self.css, title="Flux LoRA Generation") as app:
145
- title = gr.HTML(
146
- """<h1>Flux LoRA Generation</h1>""",
147
- elem_id="title",
148
- )
149
- selected_index = gr.State(None)
150
-
151
- with gr.Row():
152
- with gr.Column(scale=3):
153
- prompt = gr.Textbox(label="Prompt", lines=1, placeholder="Choose the LoRA and type the prompt")
154
- with gr.Column(scale=1, elem_id="gen_column"):
155
- generate_button = gr.Button("Generate", variant="primary", elem_id="gen_btn")
156
- with gr.Row():
157
- with gr.Column():
158
- selected_info = gr.Markdown("")
159
- gallery = gr.Gallery(
160
- [(item["image"], item["title"]) for item in self.loras],
161
- label="LoRA Collection",
162
- allow_preview=False,
163
- columns=3,
164
- elem_id="gallery",
165
- show_share_button=False
166
- )
167
- with gr.Group():
168
- custom_lora = gr.Textbox(label="Enter Custom LoRA", placeholder="prithivMLmods/Canopus-LoRA-Flux-Anime")
169
- gr.Markdown("[Check the list of FLUX LoRA's](https://huggingface.co/models?other=base_model:adapter:black-forest-labs/FLUX.1-dev)", elem_id="lora_list")
170
- custom_lora_info = gr.HTML(visible=False)
171
- custom_lora_button = gr.Button("Remove custom LoRA", visible=False)
172
- with gr.Column():
173
- progress_bar = gr.Markdown(elem_id="progress", visible=False)
174
- result = gr.Image(label="Generated Image")
175
-
176
- with gr.Row():
177
- with gr.Accordion("Advanced Settings", open=False):
178
- with gr.Row():
179
- input_image = gr.Image(label="Input image", type="filepath")
180
- image_strength = gr.Slider(label="Denoise Strength", info="Lower means more image influence", minimum=0.1, maximum=1.0, step=0.01, value=0.75)
181
- with gr.Column():
182
- with gr.Row():
183
- cfg_scale = gr.Slider(label="CFG Scale", minimum=1, maximum=20, step=0.5, value=3.5)
184
- steps = gr.Slider(label="Steps", minimum=1, maximum=50, step=1, value=28)
185
- with gr.Row():
186
- width = gr.Slider(label="Width", minimum=256, maximum=1536, step=64, value=1024)
187
- height = gr.Slider(label="Height", minimum=256, maximum=1536, step=64, value=1024)
188
- with gr.Row():
189
- randomize_seed = gr.Checkbox(True, label="Randomize seed")
190
- seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0, randomize=True)
191
- lora_scale = gr.Slider(label="LoRA Scale", minimum=0, maximum=3, step=0.01, value=0.95)
192
- # Prompt Enhancer Section
193
- with gr.Group():
194
- use_enhancer = gr.Checkbox(label="Use Prompt Enhancer", value=True)
195
- show_enhanced_prompt = gr.Checkbox(label="Display Enhanced Prompt", value=False)
196
- enhanced_prompt_box = gr.Textbox(label="Enhanced Prompt", lines=3, visible=False)
197
-
198
- gallery.select(
199
- update_selection,
200
- inputs=[width, height, gr.State(self.loras)],
201
- outputs=[prompt, selected_info, selected_index, width, height]
202
- )
203
- custom_lora.input(
204
- add_custom_lora,
205
- inputs=[custom_lora, gr.State(self.loras)],
206
- outputs=[custom_lora_info, custom_lora_button, gallery, selected_info, selected_index, prompt]
207
- )
208
- custom_lora_button.click(
209
- remove_custom_lora,
210
- outputs=[custom_lora_info, custom_lora_button, gallery, selected_info, selected_index, custom_lora]
211
- )
212
-
213
- # Toggle the visibility of the enhanced prompt textbox based on the checkbox state.
214
- show_enhanced_prompt.change(fn=lambda show: gr.update(visible=show),
215
- inputs=show_enhanced_prompt,
216
- outputs=enhanced_prompt_box)
217
-
218
- gr.on(
219
- triggers=[generate_button.click, prompt.submit],
220
- fn=self.run_lora,
221
- inputs=[prompt, input_image, image_strength, cfg_scale, steps, selected_index, randomize_seed, seed, width, height, lora_scale, use_enhancer],
222
- outputs=[result, seed, progress_bar, enhanced_prompt_box]
223
- )
224
-
225
- # Credits section added at the bottom
226
- with gr.Row():
227
- gr.HTML("<div style='text-align:center; font-size:0.9em; margin-top:20px;'>Credits: <a href='https://ruslanmv.com' target='_blank'>ruslanmv.com</a></div>")
228
-
229
- return app
230
-
231
- if __name__ == "__main__":
232
- model_manager = ModelManager()
233
- frontend = Frontend(model_manager)
234
- app = frontend.create_ui()
235
- app.queue()
236
- app.launch(debug=True)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
flux_app/frontend_v1.py DELETED
@@ -1,216 +0,0 @@
1
- # frontend.py
2
- import gradio as gr
3
- import sys
4
- import os
5
-
6
- # Add the parent directory to sys.path
7
- parent_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
8
- sys.path.insert(0, parent_dir)
9
- #print(sys.path) #DEBUG
10
-
11
- from flux_app.backend import ModelManager # Absolute import
12
- from flux_app.config import MAX_SEED # Absolute import
13
- from flux_app.lora_handling import (
14
- add_custom_lora, remove_custom_lora, prepare_prompt,
15
- unload_lora_weights, load_lora_weights_into_pipeline, update_selection
16
- )
17
- from flux_app.utilities import randomize_seed_if_needed, calculateDuration # Absolute import
18
- import spaces
19
-
20
-
21
- # Dummy loras data for initial UI setup.
22
- initial_loras = [
23
- {"image": "placeholder.jpg", "title": "Placeholder LoRA", "repo": "placeholder/repo", "weights": None, "trigger_word": ""},
24
- ]
25
-
26
- class Frontend:
27
- def __init__(self, model_manager: ModelManager):
28
- self.model_manager = model_manager
29
- self.loras = initial_loras
30
- self.load_initial_loras()
31
- self.css = self.define_css()
32
-
33
- def define_css(self):
34
- # A cleaner, professional CSS styling.
35
- return '''
36
- /* Title Styling */
37
- #title {
38
- text-align: center;
39
- margin-bottom: 20px;
40
- }
41
- #title h1 {
42
- font-size: 2.5rem;
43
- margin: 0;
44
- color: #333;
45
- }
46
- /* Button and Column Styling */
47
- #gen_btn {
48
- width: 100%;
49
- padding: 12px;
50
- font-weight: bold;
51
- border-radius: 5px;
52
- }
53
- #gen_column {
54
- display: flex;
55
- align-items: center;
56
- justify-content: center;
57
- }
58
- /* Gallery and List Styling */
59
- #gallery .grid-wrap {
60
- margin-top: 15px;
61
- }
62
- #lora_list {
63
- background-color: #f5f5f5;
64
- padding: 10px;
65
- border-radius: 4px;
66
- font-size: 0.9rem;
67
- }
68
- .card_internal {
69
- display: flex;
70
- align-items: center;
71
- height: 100px;
72
- margin-top: 10px;
73
- }
74
- .card_internal img {
75
- margin-right: 10px;
76
- }
77
- .styler {
78
- --form-gap-width: 0px !important;
79
- }
80
- /* Progress Bar Styling */
81
- .progress-container {
82
- width: 100%;
83
- height: 20px;
84
- background-color: #e0e0e0;
85
- border-radius: 10px;
86
- overflow: hidden;
87
- margin-bottom: 20px;
88
- }
89
- .progress-bar {
90
- height: 100%;
91
- background-color: #4f46e5;
92
- transition: width 0.3s ease-in-out;
93
- width: calc(var(--current) / var(--total) * 100%);
94
- }
95
- '''
96
-
97
- def load_initial_loras(self):
98
- try:
99
- from flux_app.lora import loras as loras_list # Absolute import
100
- self.loras = loras_list
101
- except ImportError:
102
- print("Warning: lora.py not found, using placeholder LoRAs.")
103
- pass
104
-
105
- @spaces.GPU(duration=100)
106
- def run_lora(self, prompt, image_input, image_strength, cfg_scale, steps, selected_index, randomize_seed, seed, width, height, lora_scale, progress=gr.Progress(track_tqdm=True)):
107
- seed = randomize_seed_if_needed(randomize_seed, seed, MAX_SEED)
108
- prompt_mash = prepare_prompt(prompt, selected_index, self.loras)
109
- selected_lora = self.loras[selected_index]
110
-
111
- unload_lora_weights(self.model_manager.pipe, self.model_manager.pipe_i2i)
112
- pipe_to_use = self.model_manager.pipe_i2i if image_input is not None else self.model_manager.pipe
113
- load_lora_weights_into_pipeline(pipe_to_use, selected_lora["repo"], selected_lora.get("weights"))
114
-
115
- if image_input is not None:
116
- final_image = self.model_manager.generate_image_to_image(
117
- prompt_mash, image_input, image_strength, steps, cfg_scale, width, height, lora_scale, seed
118
- )
119
- yield final_image, seed, gr.update(visible=False)
120
- else:
121
- image_generator = self.model_manager.generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scale)
122
- final_image = None
123
- step_counter = 0
124
- for image in image_generator:
125
- step_counter += 1
126
- final_image = image
127
- progress_bar = f'<div class="progress-container"><div class="progress-bar" style="--current: {step_counter}; --total: {steps};"></div></div>'
128
- yield image, seed, gr.update(value=progress_bar, visible=True)
129
-
130
- yield final_image, seed, gr.update(value=progress_bar, visible=False)
131
-
132
- def create_ui(self):
133
- # Using a base theme for a clean and professional look.
134
- with gr.Blocks(theme=gr.themes.Base(), css=self.css, title="Flux LoRA Generation") as app:
135
- title = gr.HTML(
136
- """<h1>Flux LoRA Generation</h1>""",
137
- elem_id="title",
138
- )
139
- selected_index = gr.State(None)
140
-
141
- with gr.Row():
142
- with gr.Column(scale=3):
143
- prompt = gr.Textbox(label="Prompt", lines=1, placeholder="Choose the LoRA and type the prompt")
144
- with gr.Column(scale=1, elem_id="gen_column"):
145
- generate_button = gr.Button("Generate", variant="primary", elem_id="gen_btn")
146
- with gr.Row():
147
- with gr.Column():
148
- selected_info = gr.Markdown("")
149
- gallery = gr.Gallery(
150
- [(item["image"], item["title"]) for item in self.loras],
151
- label="LoRA Collection",
152
- allow_preview=False,
153
- columns=3,
154
- elem_id="gallery",
155
- show_share_button=False
156
- )
157
- with gr.Group():
158
- custom_lora = gr.Textbox(label="Enter Custom LoRA", placeholder="prithivMLmods/Canopus-LoRA-Flux-Anime")
159
- gr.Markdown("[Check the list of FLUX LoRA's](https://huggingface.co/models?other=base_model:adapter:black-forest-labs/FLUX.1-dev)", elem_id="lora_list")
160
- custom_lora_info = gr.HTML(visible=False)
161
- custom_lora_button = gr.Button("Remove custom LoRA", visible=False)
162
- with gr.Column():
163
- progress_bar = gr.Markdown(elem_id="progress", visible=False)
164
- result = gr.Image(label="Generated Image")
165
-
166
- with gr.Row():
167
- with gr.Accordion("Advanced Settings", open=False):
168
- with gr.Row():
169
- input_image = gr.Image(label="Input image", type="filepath")
170
- image_strength = gr.Slider(label="Denoise Strength", info="Lower means more image influence", minimum=0.1, maximum=1.0, step=0.01, value=0.75)
171
- with gr.Column():
172
- with gr.Row():
173
- cfg_scale = gr.Slider(label="CFG Scale", minimum=1, maximum=20, step=0.5, value=3.5)
174
- steps = gr.Slider(label="Steps", minimum=1, maximum=50, step=1, value=28)
175
- with gr.Row():
176
- width = gr.Slider(label="Width", minimum=256, maximum=1536, step=64, value=1024)
177
- height = gr.Slider(label="Height", minimum=256, maximum=1536, step=64, value=1024)
178
- with gr.Row():
179
- randomize_seed = gr.Checkbox(True, label="Randomize seed")
180
- seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0, randomize=True)
181
- lora_scale = gr.Slider(label="LoRA Scale", minimum=0, maximum=3, step=0.01, value=0.95)
182
-
183
- gallery.select(
184
- update_selection,
185
- inputs=[width, height, gr.State(self.loras)],
186
- outputs=[prompt, selected_info, selected_index, width, height]
187
- )
188
- custom_lora.input(
189
- add_custom_lora,
190
- inputs=[custom_lora, gr.State(self.loras)],
191
- outputs=[custom_lora_info, custom_lora_button, gallery, selected_info, selected_index, prompt]
192
- )
193
- custom_lora_button.click(
194
- remove_custom_lora,
195
- outputs=[custom_lora_info, custom_lora_button, gallery, selected_info, selected_index, custom_lora]
196
- )
197
-
198
- gr.on(
199
- triggers=[generate_button.click, prompt.submit],
200
- fn=self.run_lora,
201
- inputs=[prompt, input_image, image_strength, cfg_scale, steps, selected_index, randomize_seed, seed, width, height, lora_scale],
202
- outputs=[result, seed, progress_bar]
203
- )
204
-
205
- # Credits section added at the bottom
206
- with gr.Row():
207
- gr.HTML("<div style='text-align:center; font-size:0.9em; margin-top:20px;'>Credits: <a href='https://ruslanmv.com' target='_blank'>ruslanmv.com</a></div>")
208
-
209
- return app
210
-
211
- if __name__ == "__main__":
212
- model_manager = ModelManager()
213
- frontend = Frontend(model_manager)
214
- app = frontend.create_ui()
215
- app.queue()
216
- app.launch()