prithivMLmods commited on
Commit
d2ba75e
·
verified ·
1 Parent(s): 862a7a5

Delete app.py

Browse files
Files changed (1) hide show
  1. app.py +0 -407
app.py DELETED
@@ -1,407 +0,0 @@
1
- #!/usr/bin/env python
2
- # Permission is hereby granted, free of charge, to any person obtaining a copy
3
- # of this software and associated documentation files (the "Software"), to deal
4
- # in the Software without restriction, including without limitation the rights
5
- # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
6
- # copies of the Software, and to permit persons to whom the Software is
7
- # furnished to do so, subject to the following conditions:
8
- #
9
- # The above copyright notice and this permission notice shall be included in
10
- # all copies or substantial portions of the Software.
11
-
12
- import os
13
- import gradio as gr
14
- import json
15
- import logging
16
- import torch
17
- from PIL import Image
18
- import spaces
19
- from diffusers import DiffusionPipeline, AutoencoderTiny, AutoencoderKL, AutoPipelineForImage2Image
20
- from live_preview_helpers import calculate_shift, retrieve_timesteps, flux_pipe_call_that_returns_an_iterable_of_images
21
- from diffusers.utils import load_image
22
- from huggingface_hub import hf_hub_download, HfFileSystem, ModelCard, snapshot_download
23
- import copy
24
- import random
25
- import time
26
- import re
27
-
28
- # Load LoRAs from JSON file
29
- with open('loras.json', 'r') as f:
30
- loras = json.load(f)
31
-
32
- # Initialize the base model
33
- dtype = torch.bfloat16
34
- device = "cuda" if torch.cuda.is_available() else "cpu"
35
- base_model = "black-forest-labs/FLUX.1-dev"
36
-
37
- taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).to(device)
38
- good_vae = AutoencoderKL.from_pretrained(base_model, subfolder="vae", torch_dtype=dtype).to(device)
39
- pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=dtype, vae=taef1).to(device)
40
- pipe_i2i = AutoPipelineForImage2Image.from_pretrained(base_model,
41
- vae=good_vae,
42
- transformer=pipe.transformer,
43
- text_encoder=pipe.text_encoder,
44
- tokenizer=pipe.tokenizer,
45
- text_encoder_2=pipe.text_encoder_2,
46
- tokenizer_2=pipe.tokenizer_2,
47
- torch_dtype=dtype
48
- )
49
-
50
- MAX_SEED = 2**32-1
51
-
52
- pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)
53
-
54
- class calculateDuration:
55
- def __init__(self, activity_name=""):
56
- self.activity_name = activity_name
57
-
58
- def __enter__(self):
59
- self.start_time = time.time()
60
- return self
61
-
62
- def __exit__(self, exc_type, exc_value, traceback):
63
- self.end_time = time.time()
64
- self.elapsed_time = self.end_time - self.start_time
65
- if self.activity_name:
66
- print(f"Elapsed time for {self.activity_name}: {self.elapsed_time:.6f} seconds")
67
- else:
68
- print(f"Elapsed time: {self.elapsed_time:.6f} seconds")
69
-
70
- def update_selection(evt: gr.SelectData, width, height):
71
- selected_lora = loras[evt.index]
72
- new_placeholder = f"Type a prompt for {selected_lora['title']}"
73
- lora_repo = selected_lora["repo"]
74
- updated_text = f"### ### Selected: [{lora_repo}](https://huggingface.co/{lora_repo}) ✅"
75
- if "aspect" in selected_lora:
76
- if selected_lora["aspect"] == "portrait":
77
- width = 768
78
- height = 1024
79
- elif selected_lora["aspect"] == "landscape":
80
- width = 1024
81
- height = 768
82
- else:
83
- width = 1024
84
- height = 1024
85
- return (
86
- gr.update(placeholder=new_placeholder),
87
- updated_text,
88
- evt.index,
89
- width,
90
- height,
91
- )
92
-
93
- @spaces.GPU(duration=100)
94
- def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scale, progress):
95
- pipe.to("cuda")
96
- generator = torch.Generator(device="cuda").manual_seed(seed)
97
- with calculateDuration("Generating image"):
98
- # Generate image
99
- for img in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
100
- prompt=prompt_mash,
101
- num_inference_steps=steps,
102
- guidance_scale=cfg_scale,
103
- width=width,
104
- height=height,
105
- generator=generator,
106
- joint_attention_kwargs={"scale": lora_scale},
107
- output_type="pil",
108
- good_vae=good_vae,
109
- ):
110
- yield img
111
-
112
- def generate_image_to_image(prompt_mash, image_input_path, image_strength, steps, cfg_scale, width, height, lora_scale, seed):
113
- generator = torch.Generator(device="cuda").manual_seed(seed)
114
- pipe_i2i.to("cuda")
115
- image_input = load_image(image_input_path)
116
- final_image = pipe_i2i(
117
- prompt=prompt_mash,
118
- image=image_input,
119
- strength=image_strength,
120
- num_inference_steps=steps,
121
- guidance_scale=cfg_scale,
122
- width=width,
123
- height=height,
124
- generator=generator,
125
- joint_attention_kwargs={"scale": lora_scale},
126
- output_type="pil",
127
- ).images[0]
128
- return final_image
129
-
130
- @spaces.GPU(duration=100)
131
- def run_lora(prompt, image_input, image_strength, cfg_scale, steps, selected_index, randomize_seed, seed, width, height, lora_scale, progress=gr.Progress(track_tqdm=True)):
132
- if selected_index is None:
133
- raise gr.Error("You must select a LoRA before proceeding.")
134
- selected_lora = loras[selected_index]
135
- lora_path = selected_lora["repo"]
136
- trigger_word = selected_lora["trigger_word"]
137
- if(trigger_word):
138
- if "trigger_position" in selected_lora:
139
- if selected_lora["trigger_position"] == "prepend":
140
- prompt_mash = f"{trigger_word} {prompt}"
141
- else:
142
- prompt_mash = f"{prompt} {trigger_word}"
143
- else:
144
- prompt_mash = f"{trigger_word} {prompt}"
145
- else:
146
- prompt_mash = prompt
147
-
148
- with calculateDuration("Unloading LoRA"):
149
- pipe.unload_lora_weights()
150
- pipe_i2i.unload_lora_weights()
151
-
152
- # Load LoRA weights
153
- with calculateDuration(f"Loading LoRA weights for {selected_lora['title']}"):
154
- pipe_to_use = pipe_i2i if image_input is not None else pipe
155
- weight_name = selected_lora.get("weights", None)
156
-
157
- pipe_to_use.load_lora_weights(
158
- lora_path,
159
- weight_name=weight_name,
160
- low_cpu_mem_usage=True
161
- )
162
-
163
- # Set random seed for reproducibility
164
- with calculateDuration("Randomizing seed"):
165
- if randomize_seed:
166
- seed = random.randint(0, MAX_SEED)
167
-
168
- if(image_input is not None):
169
-
170
- final_image = generate_image_to_image(prompt_mash, image_input, image_strength, steps, cfg_scale, width, height, lora_scale, seed)
171
- yield final_image, seed, gr.update(visible=False)
172
- else:
173
- image_generator = generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scale, progress)
174
-
175
- # Consume the generator to get the final image
176
- final_image = None
177
- step_counter = 0
178
- for image in image_generator:
179
- step_counter+=1
180
- final_image = image
181
- progress_bar = f'<div class="progress-container"><div class="progress-bar" style="--current: {step_counter}; --total: {steps};"></div></div>'
182
- yield image, seed, gr.update(value=progress_bar, visible=True)
183
-
184
- yield final_image, seed, gr.update(value=progress_bar, visible=False)
185
-
186
- def get_huggingface_safetensors(link):
187
- split_link = link.split("/")
188
- if len(split_link) != 2:
189
- raise Exception("Invalid Hugging Face repository link format.")
190
-
191
- print(f"Repository attempted: {split_link}")
192
-
193
- # Load model card
194
- model_card = ModelCard.load(link)
195
- base_model = model_card.data.get("base_model")
196
- print(f"Base model: {base_model}")
197
-
198
- # Validate model type
199
- acceptable_models = {"black-forest-labs/FLUX.1-dev", "black-forest-labs/FLUX.1-schnell"}
200
-
201
- models_to_check = base_model if isinstance(base_model, list) else [base_model]
202
-
203
- if not any(model in acceptable_models for model in models_to_check):
204
- raise Exception("Not a FLUX LoRA!")
205
-
206
- # Extract image and trigger word
207
- print("Before trying to get image")
208
- image_path = model_card.data.get("widget", [{}])[0].get("output", {}).get("url", None)
209
- print(f"Image path {image_path}")
210
- trigger_word = model_card.data.get("instance_prompt", "")
211
- print(f"Image path {trigger_word}")
212
- image_url = f"https://huggingface.co/{link}/resolve/main/{image_path}" if image_path else None
213
- print(f"Image URL {image_url}")
214
-
215
-
216
- # Initialize Hugging Face file system
217
- fs = HfFileSystem()
218
- try:
219
- list_of_files = fs.ls(link, detail=False)
220
-
221
- # Initialize variables for safetensors selection
222
- safetensors_name = None
223
- highest_trained_file = None
224
- highest_steps = -1
225
- last_safetensors_file = None
226
- step_pattern = re.compile(r"_0{3,}\d+") # Detects step count `_000...`
227
-
228
- for file in list_of_files:
229
- filename = file.split("/")[-1]
230
-
231
- # Select safetensors file
232
- if filename.endswith(".safetensors"):
233
- last_safetensors_file = filename # Track last encountered file
234
-
235
- match = step_pattern.search(filename)
236
- if not match:
237
- # Found a full model without step numbers, return immediately
238
- safetensors_name = filename
239
- break
240
- else:
241
- # Extract step count and track highest
242
- steps = int(match.group().lstrip("_"))
243
- if steps > highest_steps:
244
- highest_trained_file = filename
245
- highest_steps = steps
246
-
247
- # Select an image file if not found in model card
248
- if not image_url and filename.lower().endswith((".jpg", ".jpeg", ".png", ".webp")):
249
- image_url = f"https://huggingface.co/{link}/resolve/main/{filename}"
250
-
251
- # If no full model found, fall back to the most trained safetensors file
252
- if not safetensors_name:
253
- safetensors_name = highest_trained_file if highest_trained_file else last_safetensors_file
254
-
255
- # If still no safetensors file found, raise an exception
256
- if not safetensors_name:
257
- raise Exception("No valid *.safetensors file found in the repository.")
258
-
259
- except Exception as e:
260
- print(e)
261
- raise Exception("You didn't include a valid Hugging Face repository with a *.safetensors LoRA")
262
-
263
- return split_link[1], link, safetensors_name, trigger_word, image_url
264
-
265
-
266
- def check_custom_model(link):
267
- print(f"Checking a custom model on: {link}")
268
- if(link.startswith("https://")):
269
- if(link.startswith("https://huggingface.co") or link.startswith("https://www.huggingface.co")):
270
- link_split = link.split("huggingface.co/")
271
- return get_huggingface_safetensors(link_split[1])
272
- else:
273
- return get_huggingface_safetensors(link)
274
-
275
- def add_custom_lora(custom_lora):
276
- global loras
277
- if(custom_lora):
278
- try:
279
- title, repo, path, trigger_word, image = check_custom_model(custom_lora)
280
- print(f"Loaded custom LoRA: {repo}")
281
- card = f'''
282
- <div class="custom_lora_card">
283
- <span>Loaded custom LoRA:</span>
284
- <div class="card_internal">
285
- <img src="{image}" />
286
- <div>
287
- <h3>{title}</h3>
288
- <small>{"Using: <code><b>"+trigger_word+"</code></b> as the trigger word" if trigger_word else "No trigger word found. If there's a trigger word, include it in your prompt"}<br></small>
289
- </div>
290
- </div>
291
- </div>
292
- '''
293
- existing_item_index = next((index for (index, item) in enumerate(loras) if item['repo'] == repo), None)
294
- if(not existing_item_index):
295
- new_item = {
296
- "image": image,
297
- "title": title,
298
- "repo": repo,
299
- "weights": path,
300
- "trigger_word": trigger_word
301
- }
302
- print(new_item)
303
- existing_item_index = len(loras)
304
- loras.append(new_item)
305
-
306
- return gr.update(visible=True, value=card), gr.update(visible=True), gr.Gallery(selected_index=None), f"Custom: {path}", existing_item_index, trigger_word
307
- except Exception as e:
308
- gr.Warning(f"Invalid LoRA: either you entered an invalid link, or a non-FLUX LoRA, this was the issue: {e}")
309
- return gr.update(visible=True, value=f"Invalid LoRA: either you entered an invalid link, a non-FLUX LoRA"), gr.update(visible=True), gr.update(), "", None, ""
310
- else:
311
- return gr.update(visible=False), gr.update(visible=False), gr.update(), "", None, ""
312
-
313
- def remove_custom_lora():
314
- return gr.update(visible=False), gr.update(visible=False), gr.update(), "", None, ""
315
-
316
- run_lora.zerogpu = True
317
-
318
- css = '''
319
- #gen_btn{height: 100%}
320
- #gen_column{align-self: stretch}
321
- #title{text-align: center}
322
- #title h1{font-size: 3em; display:inline-flex; align-items:center}
323
- #title img{width: 100px; margin-right: 0.5em}
324
- #gallery .grid-wrap{height: 10vh}
325
- #lora_list{background: var(--block-background-fill);padding: 0 1em .3em; font-size: 90%}
326
- .card_internal{display: flex;height: 100px;margin-top: .5em}
327
- .card_internal img{margin-right: 1em}
328
- .styler{--form-gap-width: 0px !important}
329
- #progress{height:30px}
330
- #progress .generating{display:none}
331
- .progress-container {width: 100%;height: 30px;background-color: #f0f0f0;border-radius: 15px;overflow: hidden;margin-bottom: 20px}
332
- .progress-bar {height: 100%;background-color: #4f46e5;width: calc(var(--current) / var(--total) * 100%);transition: width 0.5s ease-in-out}
333
- '''
334
-
335
- with gr.Blocks(theme="bethecloud/storj_theme", css=css, delete_cache=(60, 60)) as app:
336
- title = gr.HTML(
337
- """<h1>FLUX LoRA DLC🥳</h1>""",
338
- elem_id="title",
339
- )
340
- selected_index = gr.State(None)
341
- with gr.Row():
342
- with gr.Column(scale=3):
343
- prompt = gr.Textbox(label="Prompt", lines=1, placeholder=":/ choose the LoRA and type the prompt ")
344
- with gr.Column(scale=1, elem_id="gen_column"):
345
- generate_button = gr.Button("Generate", variant="primary", elem_id="gen_btn")
346
- with gr.Row():
347
- with gr.Column():
348
- selected_info = gr.Markdown("")
349
- gallery = gr.Gallery(
350
- [(item["image"], item["title"]) for item in loras],
351
- label="LoRA Gallery",
352
- allow_preview=False,
353
- columns=3,
354
- elem_id="gallery",
355
- show_share_button=False
356
- )
357
- with gr.Group():
358
- custom_lora = gr.Textbox(label="Custom LoRA", placeholder="prithivMLmods/Canopus-LoRA-Flux-Anime")
359
- 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")
360
- custom_lora_info = gr.HTML(visible=False)
361
- custom_lora_button = gr.Button("Remove custom LoRA", visible=False)
362
- with gr.Column():
363
- progress_bar = gr.Markdown(elem_id="progress",visible=False)
364
- result = gr.Image(label="Generated Image")
365
-
366
- with gr.Row():
367
- with gr.Accordion("Advanced Settings", open=False):
368
- with gr.Row():
369
- input_image = gr.Image(label="Input image", type="filepath")
370
- 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)
371
- with gr.Column():
372
- with gr.Row():
373
- cfg_scale = gr.Slider(label="CFG Scale", minimum=1, maximum=20, step=0.5, value=3.5)
374
- steps = gr.Slider(label="Steps", minimum=1, maximum=50, step=1, value=28)
375
-
376
- with gr.Row():
377
- width = gr.Slider(label="Width", minimum=256, maximum=1536, step=64, value=1024)
378
- height = gr.Slider(label="Height", minimum=256, maximum=1536, step=64, value=1024)
379
-
380
- with gr.Row():
381
- randomize_seed = gr.Checkbox(True, label="Randomize seed")
382
- seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0, randomize=True)
383
- lora_scale = gr.Slider(label="LoRA Scale", minimum=0, maximum=3, step=0.01, value=0.95)
384
-
385
- gallery.select(
386
- update_selection,
387
- inputs=[width, height],
388
- outputs=[prompt, selected_info, selected_index, width, height]
389
- )
390
- custom_lora.input(
391
- add_custom_lora,
392
- inputs=[custom_lora],
393
- outputs=[custom_lora_info, custom_lora_button, gallery, selected_info, selected_index, prompt]
394
- )
395
- custom_lora_button.click(
396
- remove_custom_lora,
397
- outputs=[custom_lora_info, custom_lora_button, gallery, selected_info, selected_index, custom_lora]
398
- )
399
- gr.on(
400
- triggers=[generate_button.click, prompt.submit],
401
- fn=run_lora,
402
- inputs=[prompt, input_image, image_strength, cfg_scale, steps, selected_index, randomize_seed, seed, width, height, lora_scale],
403
- outputs=[result, seed, progress_bar]
404
- )
405
-
406
- app.queue()
407
- app.launch()