AlekseyCalvin commited on
Commit
2cbe354
·
verified ·
1 Parent(s): 7a7114f

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/hyper-flux1-dev-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=good_vae).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"> Novorealist LoRa-stocked Birthweek-inspired Img Manufactory for Dunova, Dunovas, & Dunovaists!</div>"""
140
+ )
141
+
142
+ # Info blob stating what the app is running
143
+ info_blob = gr.HTML(
144
+ """<div id="info_blob">Trigger LoRAs by Pre-phrasing Prompts w/: 1-5. ADU person (/'ADU woman') photo |6-15. HST style |16. how2draw |17-20.HST |21. HST Austin Osman Spare style |22. RCA |23. propaganda poster |24. SOTS art |25. pficonics |26. wh3r3sw4ld0 |27. vintage cover |28. crisp photo |29. retrofuturism |30. Film Photo |31. TOK hybrid |32. 2004 photo |33. Unexpected photo |34. flmft |35. TOK portra |36. Yearbook photo |37. Akhmatova |38. Tsvetaeva |39. Blok |40. LEN Lenin |41. Trotsky |42. Rosa Luxemburg </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=0, maximum=20, step=0.5, value=1.0)
171
+ steps = gr.Slider(label="Steps", minimum=1, maximum=50, step=1, value=12)
172
+
173
+ with gr.Row():
174
+ width = gr.Slider(label="Width", minimum=256, maximum=1536, step=64, value=1152)
175
+ height = gr.Slider(label="Height", minimum=256, maximum=1536, step=64, value=1088)
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=1.12)
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()