AlekseyCalvin commited on
Commit
917d2f3
1 Parent(s): 425be2b

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +245 -0
app.py ADDED
@@ -0,0 +1,245 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import gradio as gr
3
+ import json
4
+ import logging
5
+ import torch
6
+ from PIL import Image
7
+ from os import path
8
+ import spaces
9
+ from diffusers import DiffusionPipeline, AutoencoderTiny, AutoencoderKL
10
+ from diffusers.models.transformers import FluxTransformer2DModel
11
+ from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
12
+ from transformers import CLIPModel, CLIPProcessor, CLIPTextModel, CLIPTokenizer, CLIPConfig, T5EncoderModel, T5Tokenizer
13
+ import copy
14
+ import random
15
+ import time
16
+ from live_preview_helpers import calculate_shift, retrieve_timesteps, flux_pipe_call_that_returns_an_iterable_of_images
17
+ from huggingface_hub import HfFileSystem, ModelCard
18
+ from huggingface_hub import login, hf_hub_download
19
+ import safetensors.torch
20
+ from safetensors.torch import load_file
21
+ hf_token = os.environ.get("HF_TOKEN")
22
+ login(token=hf_token)
23
+
24
+ cache_path = path.join(path.dirname(path.abspath(__file__)), "models")
25
+ os.environ["TRANSFORMERS_CACHE"] = cache_path
26
+ os.environ["HF_HUB_CACHE"] = cache_path
27
+ os.environ["HF_HOME"] = cache_path
28
+
29
+ torch.set_float32_matmul_precision("medium")
30
+
31
+ #torch._inductor.config.conv_1x1_as_mm = True
32
+ #torch._inductor.config.coordinate_descent_tuning = True
33
+ #torch._inductor.config.epilogue_fusion = False
34
+ #torch._inductor.config.coordinate_descent_check_all_directions = False
35
+
36
+ # Load LoRAs from JSON file
37
+ with open('loras.json', 'r') as f:
38
+ loras = json.load(f)
39
+
40
+ # Initialize the base model
41
+ dtype = torch.bfloat16
42
+ device = "cuda" if torch.cuda.is_available() else "cpu"
43
+
44
+ taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).to(device)
45
+ good_vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=dtype).to(device)
46
+
47
+ #pipe = DiffusionPipeline.from_pretrained(
48
+ #"AlekseyCalvin/Flux_Dedistilled_Mix_byWikeeyang_Diffusers",
49
+ #custom_pipeline="jimmycarter/LibreFLUX",
50
+ #use_safetensors=True,
51
+ #torch_dtype=torch.bfloat16,
52
+ #trust_remote_code=True,
53
+ #).to(device)
54
+
55
+ pipe = FluxPipeline.from_pretrained("AlekseyCalvin/FluxFusionV2_by_Anibaal_Diffusers", torch_dtype=torch.bfloat16)
56
+ pipe.to(device="cuda", dtype=torch.bfloat16)
57
+
58
+ clipmodel = 'norm'
59
+ if clipmodel == "long":
60
+ model_id = "zer0int/LongCLIP-GmP-ViT-L-14"
61
+ config = CLIPConfig.from_pretrained(model_id)
62
+ maxtokens = 77
63
+ if clipmodel == "norm":
64
+ model_id = "zer0int/CLIP-GmP-ViT-L-14"
65
+ config = CLIPConfig.from_pretrained(model_id)
66
+ maxtokens = 77
67
+ clip_model = CLIPModel.from_pretrained(model_id, torch_dtype=torch.bfloat16, config=config, ignore_mismatched_sizes=True).to("cuda")
68
+ clip_processor = CLIPProcessor.from_pretrained(model_id, padding="max_length", max_length=maxtokens, ignore_mismatched_sizes=True, return_tensors="pt", truncation=True)
69
+
70
+ pipe.tokenizer = clip_processor.tokenizer
71
+ pipe.text_encoder = clip_model.text_model
72
+ pipe.tokenizer_max_length = maxtokens
73
+ pipe.text_encoder.dtype = torch.bfloat16
74
+ pipe.vae = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).to("cuda")
75
+
76
+ # Load LoRAs from JSON file
77
+ with open('loras.json', 'r') as f:
78
+ loras = json.load(f)
79
+
80
+ MAX_SEED = 2**32-1
81
+
82
+ class calculateDuration:
83
+ def __init__(self, activity_name=""):
84
+ self.activity_name = activity_name
85
+
86
+ def __enter__(self):
87
+ self.start_time = time.time()
88
+ return self
89
+
90
+ def __exit__(self, exc_type, exc_value, traceback):
91
+ self.end_time = time.time()
92
+ self.elapsed_time = self.end_time - self.start_time
93
+ if self.activity_name:
94
+ print(f"Elapsed time for {self.activity_name}: {self.elapsed_time:.6f} seconds")
95
+ else:
96
+ print(f"Elapsed time: {self.elapsed_time:.6f} seconds")
97
+
98
+
99
+ def update_selection(evt: gr.SelectData, width, height):
100
+ selected_lora = loras[evt.index]
101
+ new_placeholder = f"Type a prompt for {selected_lora['title']}"
102
+ lora_repo = selected_lora["repo"]
103
+ updated_text = f"### Selected: [{lora_repo}](https://huggingface.co/{lora_repo}) ✨"
104
+ if "aspect" in selected_lora:
105
+ if selected_lora["aspect"] == "portrait":
106
+ width = 768
107
+ height = 1024
108
+ elif selected_lora["aspect"] == "landscape":
109
+ width = 1024
110
+ height = 768
111
+ return (
112
+ gr.update(placeholder=new_placeholder),
113
+ updated_text,
114
+ evt.index,
115
+ width,
116
+ height,
117
+ )
118
+
119
+ @spaces.GPU(duration=70)
120
+ def generate_image(prompt, trigger_word, steps, seed, cfg_scale, width, height, lora_scale, progress):
121
+ pipe.to("cuda")
122
+ generator = torch.Generator(device="cuda").manual_seed(seed)
123
+
124
+ with calculateDuration("Generating image"):
125
+ # Generate image
126
+ image = pipe(
127
+ prompt=f"{prompt} {trigger_word}",
128
+ num_inference_steps=steps,
129
+ guidance_scale=cfg_scale,
130
+ width=width,
131
+ height=height,
132
+ generator=generator,
133
+ joint_attention_kwargs={"scale": lora_scale},
134
+ ).images[0]
135
+ return image
136
+
137
+ def run_lora(prompt, cfg_scale, steps, selected_index, randomize_seed, seed, width, height, lora_scale, progress=gr.Progress(track_tqdm=True)):
138
+ if selected_index is None:
139
+ raise gr.Error("You must select a LoRA before proceeding.")
140
+
141
+ selected_lora = loras[selected_index]
142
+ lora_path = selected_lora["repo"]
143
+ trigger_word = selected_lora["trigger_word"]
144
+ if(trigger_word):
145
+ if "trigger_position" in selected_lora:
146
+ if selected_lora["trigger_position"] == "prepend":
147
+ prompt_mash = f"{trigger_word} {prompt}"
148
+ else:
149
+ prompt_mash = f"{prompt} {trigger_word}"
150
+ else:
151
+ prompt_mash = f"{trigger_word} {prompt}"
152
+ else:
153
+ prompt_mash = prompt
154
+
155
+ # Load LoRA weights
156
+ with calculateDuration(f"Loading LoRA weights for {selected_lora['title']}"):
157
+ if "weights" in selected_lora:
158
+ pipe.load_lora_weights(lora_path, weight_name=selected_lora["weights"])
159
+ else:
160
+ pipe.load_lora_weights(lora_path)
161
+
162
+ # Set random seed for reproducibility
163
+ with calculateDuration("Randomizing seed"):
164
+ if randomize_seed:
165
+ seed = random.randint(0, MAX_SEED)
166
+
167
+ image = generate_image(prompt, trigger_word, steps, seed, cfg_scale, width, height, lora_scale, progress)
168
+ pipe.to("cpu")
169
+ pipe.unload_lora_weights()
170
+ return image, seed
171
+
172
+ run_lora.zerogpu = True
173
+
174
+ css = '''
175
+ #gen_btn{height: 100%}
176
+ #title{text-align: center}
177
+ #title h1{font-size: 3em; display:inline-flex; align-items:center}
178
+ #title img{width: 100px; margin-right: 0.5em}
179
+ #gallery .grid-wrap{height: 10vh}
180
+ '''
181
+ with gr.Blocks(theme=gr.themes.Soft(), css=css) as app:
182
+ title = gr.HTML(
183
+ """<h1><img src="https://huggingface.co/spaces/multimodalart/flux-lora-the-explorer/resolve/main/flux_lora.png" alt="LoRA"> SOONfactory </h1>""",
184
+ elem_id="title",
185
+ )
186
+ # Info blob stating what the app is running
187
+ info_blob = gr.HTML(
188
+ """<div id="info_blob"> Activist & Futurealist Fast FLUX De-Distilled LoRA Gallery (Running on the the FLUX De-Distilled Mix Tuned V1 by wikeeyang)</div>"""
189
+ )
190
+
191
+ # Info blob stating what the app is running
192
+ info_blob = gr.HTML(
193
+ """<div id="info_blob">Prephrase prompts w/: 1.RCA style 2. HST style autochrome 3. HST style 4.TOK hybrid 5.2004 photo 6.HST style 7.LEN Vladimir Lenin 8.TOK portra 9.HST portrait 10.flmft 11.HST in Peterhof 12.HST Soviet kodachrome 13. SOTS art 14.HST 15.photo 16.pficonics 17.wh3r3sw4ld0 18.retrofuturism 19-24.HST style photo 25.vintage cover </div>"""
194
+ )
195
+ selected_index = gr.State(None)
196
+ with gr.Row():
197
+ with gr.Column(scale=3):
198
+ prompt = gr.Textbox(label="Prompt", lines=1, placeholder="Select LoRa/Style & type prompt!")
199
+ with gr.Column(scale=1, elem_id="gen_column"):
200
+ generate_button = gr.Button("Generate", variant="primary", elem_id="gen_btn")
201
+ with gr.Row():
202
+ with gr.Column(scale=3):
203
+ selected_info = gr.Markdown("")
204
+ gallery = gr.Gallery(
205
+ [(item["image"], item["title"]) for item in loras],
206
+ label="LoRA Inventory",
207
+ allow_preview=False,
208
+ columns=3,
209
+ elem_id="gallery"
210
+ )
211
+
212
+ with gr.Column(scale=4):
213
+ result = gr.Image(label="Generated Image")
214
+
215
+ with gr.Row():
216
+ with gr.Accordion("Advanced Settings", open=True):
217
+ with gr.Column():
218
+ with gr.Row():
219
+ cfg_scale = gr.Slider(label="CFG Scale", minimum=1, maximum=20, step=0.5, value=3.0)
220
+ steps = gr.Slider(label="Steps", minimum=1, maximum=50, step=1, value=10)
221
+
222
+ with gr.Row():
223
+ width = gr.Slider(label="Width", minimum=256, maximum=1536, step=64, value=1024)
224
+ height = gr.Slider(label="Height", minimum=256, maximum=1536, step=64, value=1024)
225
+
226
+ with gr.Row():
227
+ randomize_seed = gr.Checkbox(True, label="Randomize seed")
228
+ seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0, randomize=True)
229
+ lora_scale = gr.Slider(label="LoRA Scale", minimum=0, maximum=2.0, step=0.01, value=1.0)
230
+
231
+ gallery.select(
232
+ update_selection,
233
+ inputs=[width, height],
234
+ outputs=[prompt, selected_info, selected_index, width, height]
235
+ )
236
+
237
+ gr.on(
238
+ triggers=[generate_button.click, prompt.submit],
239
+ fn=run_lora,
240
+ inputs=[prompt, cfg_scale, steps, selected_index, randomize_seed, seed, width, height, lora_scale],
241
+ outputs=[result, seed]
242
+ )
243
+
244
+ app.queue(default_concurrency_limit=2).launch(show_error=True)
245
+ app.launch()