AlekseyCalvin commited on
Commit
c3b1369
1 Parent(s): 67d51bc

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +196 -0
app.py ADDED
@@ -0,0 +1,196 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import gradio as gr
3
+ import json
4
+ import logging
5
+ import torch
6
+ from PIL import Image
7
+ import spaces
8
+ from diffusers import DiffusionPipeline, AutoencoderTiny, AutoencoderKL
9
+ import copy
10
+ import random
11
+ import time
12
+ from live_preview_helpers import calculate_shift, retrieve_timesteps, flux_pipe_call_that_returns_an_iterable_of_images
13
+ from huggingface_hub import HfFileSystem, ModelCard
14
+
15
+ from huggingface_hub import login
16
+ hf_token = os.environ.get("HF_TOKEN")
17
+ login(token=hf_token)
18
+
19
+ # Load LoRAs from JSON file
20
+ with open('loras.json', 'r') as f:
21
+ loras = json.load(f)
22
+
23
+ # Initialize the base model
24
+ dtype = torch.bfloat16
25
+ device = "cuda" if torch.cuda.is_available() else "cpu"
26
+ base_model = "John6666/real-flux-10b-schnell-fp8-flux"
27
+ taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).to(device)
28
+ good_vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=dtype).to(device)
29
+ pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=dtype, vae=taef1).to(device)
30
+
31
+ MAX_SEED = 2**32-1
32
+
33
+ class calculateDuration:
34
+ def __init__(self, activity_name=""):
35
+ self.activity_name = activity_name
36
+
37
+ def __enter__(self):
38
+ self.start_time = time.time()
39
+ return self
40
+
41
+ def __exit__(self, exc_type, exc_value, traceback):
42
+ self.end_time = time.time()
43
+ self.elapsed_time = self.end_time - self.start_time
44
+ if self.activity_name:
45
+ print(f"Elapsed time for {self.activity_name}: {self.elapsed_time:.6f} seconds")
46
+ else:
47
+ print(f"Elapsed time: {self.elapsed_time:.6f} seconds")
48
+
49
+
50
+ def update_selection(evt: gr.SelectData, width, height):
51
+ selected_lora = loras[evt.index]
52
+ new_placeholder = f"Type a prompt for {selected_lora['title']}"
53
+ lora_repo = selected_lora["repo"]
54
+ updated_text = f"### Selected: [{lora_repo}](https://huggingface.co/{lora_repo}) ✨"
55
+ if "aspect" in selected_lora:
56
+ if selected_lora["aspect"] == "portrait":
57
+ width = 768
58
+ height = 1024
59
+ elif selected_lora["aspect"] == "landscape":
60
+ width = 1024
61
+ height = 768
62
+ return (
63
+ gr.update(placeholder=new_placeholder),
64
+ updated_text,
65
+ evt.index,
66
+ width,
67
+ height,
68
+ )
69
+
70
+ @spaces.GPU(duration=70)
71
+ def generate_image(prompt, trigger_word, steps, seed, cfg_scale, width, height, lora_scale, progress):
72
+ pipe.to("cuda")
73
+ generator = torch.Generator(device="cuda").manual_seed(seed)
74
+
75
+ with calculateDuration("Generating image"):
76
+ # Generate image
77
+ image = pipe(
78
+ prompt=f"{prompt} {trigger_word}",
79
+ num_inference_steps=steps,
80
+ guidance_scale=cfg_scale,
81
+ width=width,
82
+ height=height,
83
+ generator=generator,
84
+ joint_attention_kwargs={"scale": lora_scale},
85
+ ).images[0]
86
+ return image
87
+
88
+ def run_lora(prompt, cfg_scale, steps, selected_index, randomize_seed, seed, width, height, lora_scale, progress=gr.Progress(track_tqdm=True)):
89
+ if selected_index is None:
90
+ raise gr.Error("You must select a LoRA before proceeding.")
91
+
92
+ selected_lora = loras[selected_index]
93
+ lora_path = selected_lora["repo"]
94
+ trigger_word = selected_lora["trigger_word"]
95
+ if(trigger_word):
96
+ if "trigger_position" in selected_lora:
97
+ if selected_lora["trigger_position"] == "prepend":
98
+ prompt_mash = f"{trigger_word} {prompt}"
99
+ else:
100
+ prompt_mash = f"{prompt} {trigger_word}"
101
+ else:
102
+ prompt_mash = f"{trigger_word} {prompt}"
103
+ else:
104
+ prompt_mash = prompt
105
+
106
+ # Load LoRA weights
107
+ with calculateDuration(f"Loading LoRA weights for {selected_lora['title']}"):
108
+ if "weights" in selected_lora:
109
+ pipe.load_lora_weights(lora_path, weight_name=selected_lora["weights"])
110
+ else:
111
+ pipe.load_lora_weights(lora_path)
112
+
113
+ # Set random seed for reproducibility
114
+ with calculateDuration("Randomizing seed"):
115
+ if randomize_seed:
116
+ seed = random.randint(0, MAX_SEED)
117
+
118
+ image = generate_image(prompt, trigger_word, steps, seed, cfg_scale, width, height, lora_scale, progress)
119
+ pipe.to("cpu")
120
+ pipe.unload_lora_weights()
121
+ return image, seed
122
+
123
+ run_lora.zerogpu = True
124
+
125
+ css = '''
126
+ #gen_btn{height: 100%}
127
+ #title{text-align: center}
128
+ #title h1{font-size: 3em; display:inline-flex; align-items:center}
129
+ #title img{width: 100px; margin-right: 0.5em}
130
+ #gallery .grid-wrap{height: 10vh}
131
+ '''
132
+ with gr.Blocks(theme=gr.themes.Soft(), css=css) as app:
133
+ title = gr.HTML(
134
+ """<h1><img src="https://huggingface.co/spaces/multimodalart/flux-lora-the-explorer/resolve/main/flux_lora.png" alt="LoRA"> SOONfactory </h1>""",
135
+ elem_id="title",
136
+ )
137
+ # Info blob stating what the app is running
138
+ info_blob = gr.HTML(
139
+ """<div id="info_blob"> Activist & Futurealist LoRa-stocked Img Manufactory (currently on the Acorn is Spinning Schnell V1 model checkpoint, 4-step, by Seeker70) )</div>"""
140
+ )
141
+
142
+ # Info blob stating what the app is running
143
+ info_blob = gr.HTML(
144
+ """<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>"""
145
+ )
146
+ selected_index = gr.State(None)
147
+ with gr.Row():
148
+ with gr.Column(scale=3):
149
+ prompt = gr.Textbox(label="Prompt", lines=1, placeholder="Select LoRa/Style & type prompt!")
150
+ with gr.Column(scale=1, elem_id="gen_column"):
151
+ generate_button = gr.Button("Generate", variant="primary", elem_id="gen_btn")
152
+ with gr.Row():
153
+ with gr.Column(scale=3):
154
+ selected_info = gr.Markdown("")
155
+ gallery = gr.Gallery(
156
+ [(item["image"], item["title"]) for item in loras],
157
+ label="LoRA Inventory",
158
+ allow_preview=False,
159
+ columns=3,
160
+ elem_id="gallery"
161
+ )
162
+
163
+ with gr.Column(scale=4):
164
+ result = gr.Image(label="Generated Image")
165
+
166
+ with gr.Row():
167
+ with gr.Accordion("Advanced Settings", open=True):
168
+ with gr.Column():
169
+ with gr.Row():
170
+ cfg_scale = gr.Slider(label="CFG Scale", minimum=1, maximum=20, step=0.5, value=1.0)
171
+ steps = gr.Slider(label="Steps", minimum=1, maximum=50, step=1, value=5)
172
+
173
+ with gr.Row():
174
+ width = gr.Slider(label="Width", minimum=256, maximum=1536, step=64, value=768)
175
+ height = gr.Slider(label="Height", minimum=256, maximum=1536, step=64, value=1024)
176
+
177
+ with gr.Row():
178
+ randomize_seed = gr.Checkbox(True, label="Randomize seed")
179
+ seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0, randomize=True)
180
+ lora_scale = gr.Slider(label="LoRA Scale", minimum=0, maximum=2.0, step=0.01, value=0.6)
181
+
182
+ gallery.select(
183
+ update_selection,
184
+ inputs=[width, height],
185
+ outputs=[prompt, selected_info, selected_index, width, height]
186
+ )
187
+
188
+ gr.on(
189
+ triggers=[generate_button.click, prompt.submit],
190
+ fn=run_lora,
191
+ inputs=[prompt, cfg_scale, steps, selected_index, randomize_seed, seed, width, height, lora_scale],
192
+ outputs=[result, seed]
193
+ )
194
+
195
+ app.queue(default_concurrency_limit=2).launch(show_error=True)
196
+ app.launch()