Spaces:
Sleeping
Sleeping
AlekseyCalvin
commited on
Create app.py
Browse files
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()
|