AlekseyCalvin commited on
Commit
61321c8
·
verified ·
1 Parent(s): 4bd2f77

Upload app.py

Browse files
Files changed (1) hide show
  1. app.py +216 -0
app.py ADDED
@@ -0,0 +1,216 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import json
3
+ import logging
4
+ import argparse
5
+ import torch
6
+ import os
7
+ from os import path
8
+ from PIL import Image
9
+ import numpy as np
10
+ import spaces
11
+ import copy
12
+ import random
13
+ import time
14
+ from typing import Any, Dict, List, Optional, Union
15
+ from huggingface_hub import hf_hub_download
16
+ from diffusers import DiffusionPipeline, AutoencoderTiny, AutoPipelineForImage2Image
17
+ import safetensors.torch
18
+ from safetensors.torch import load_file
19
+ from pipeline import FluxWithCFGPipeline
20
+ from transformers import CLIPModel, CLIPProcessor, CLIPConfig
21
+ import gc
22
+ import warnings
23
+
24
+ os.environ["TRANSFORMERS_CACHE"] = cache_path
25
+ os.environ["HF_HUB_CACHE"] = cache_path
26
+ os.environ["HF_HOME"] = cache_path
27
+
28
+ device = "cuda" if torch.cuda.is_available() else "cpu"
29
+
30
+ torch.backends.cuda.matmul.allow_tf32 = True
31
+
32
+ # Load LoRAs from JSON file
33
+ with open('loras.json', 'r') as f:
34
+ loras = json.load(f)
35
+
36
+ dtype = torch.bfloat16
37
+ pipe = FluxWithCFGPipeline.from_pretrained("ostris/OpenFLUX.1", torch_dtype=dtype, text_encoder_3=None, tokenizer_3=None
38
+ ).to("cuda")
39
+ pipe.vae = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).to("cuda")
40
+
41
+ pipe.to("cuda")
42
+ clipmodel = 'norm'
43
+ if clipmodel == "long":
44
+ model_id = "zer0int/LongCLIP-GmP-ViT-L-14"
45
+ config = CLIPConfig.from_pretrained(model_id)
46
+ maxtokens = 77
47
+ if clipmodel == "norm":
48
+ model_id = "zer0int/CLIP-GmP-ViT-L-14"
49
+ config = CLIPConfig.from_pretrained(model_id)
50
+ maxtokens = 77
51
+ clip_model = CLIPModel.from_pretrained(model_id, torch_dtype=torch.bfloat16, config=config, ignore_mismatched_sizes=True).to("cuda")
52
+ clip_processor = CLIPProcessor.from_pretrained(model_id, padding="max_length", max_length=maxtokens, ignore_mismatched_sizes=True, return_tensors="pt", truncation=True)
53
+ config.text_config.max_position_embeddings = 77
54
+
55
+ pipe.tokenizer = clip_processor.tokenizer
56
+ pipe.text_encoder = clip_model.text_model
57
+ pipe.tokenizer_max_length = maxtokens
58
+ pipe.text_encoder.dtype = torch.bfloat16
59
+ torch.cuda.empty_cache()
60
+
61
+ MAX_SEED = 2**32-1
62
+
63
+ class calculateDuration:
64
+ def __init__(self, activity_name=""):
65
+ self.activity_name = activity_name
66
+
67
+ def __enter__(self):
68
+ self.start_time = time.time()
69
+ return self
70
+
71
+ def __exit__(self, exc_type, exc_value, traceback):
72
+ self.end_time = time.time()
73
+ self.elapsed_time = self.end_time - self.start_time
74
+ if self.activity_name:
75
+ print(f"Elapsed time for {self.activity_name}: {self.elapsed_time:.6f} seconds")
76
+ else:
77
+ print(f"Elapsed time: {self.elapsed_time:.6f} seconds")
78
+
79
+
80
+ def update_selection(evt: gr.SelectData, width, height):
81
+ selected_lora = loras[evt.index]
82
+ new_placeholder = f"Type a prompt for {selected_lora['title']}"
83
+ lora_repo = selected_lora["repo"]
84
+ updated_text = f"### Selected: [{lora_repo}](https://huggingface.co/{lora_repo}) ✨"
85
+ if "aspect" in selected_lora:
86
+ if selected_lora["aspect"] == "portrait":
87
+ width = 768
88
+ height = 1024
89
+ elif selected_lora["aspect"] == "landscape":
90
+ width = 1024
91
+ height = 768
92
+ return (
93
+ gr.update(placeholder=new_placeholder),
94
+ updated_text,
95
+ evt.index,
96
+ width,
97
+ height,
98
+ )
99
+
100
+ @spaces.GPU(duration=70)
101
+ def generate_image(prompt, trigger_word, steps, seed, cfg_scale, width, height, lora_scale, progress):
102
+ pipe.to("cuda")
103
+ generator = torch.Generator(device="cuda").manual_seed(seed)
104
+
105
+ with calculateDuration("Generating image"):
106
+ # Generate image
107
+ image = pipe(
108
+ prompt=f"{prompt} {trigger_word}",
109
+ num_inference_steps=steps,
110
+ guidance_scale=cfg_scale,
111
+ width=width,
112
+ height=height,
113
+ generator=generator,
114
+ joint_attention_kwargs={"scale": lora_scale},
115
+ ).images[0]
116
+ return image
117
+
118
+ def run_lora(prompt, cfg_scale, steps, selected_index, randomize_seed, seed, width, height, lora_scale, progress=gr.Progress(track_tqdm=True)):
119
+ if selected_index is None:
120
+ raise gr.Error("You must select a LoRA before proceeding.")
121
+
122
+ selected_lora = loras[selected_index]
123
+ lora_path = selected_lora["repo"]
124
+ trigger_word = selected_lora["trigger_word"]
125
+
126
+ # Load LoRA weights
127
+ with calculateDuration(f"Loading LoRA weights for {selected_lora['title']}"):
128
+ if "weights" in selected_lora:
129
+ pipe.load_lora_weights(lora_path, weight_name=selected_lora["weights"])
130
+ else:
131
+ pipe.load_lora_weights(lora_path)
132
+
133
+ # Set random seed for reproducibility
134
+ with calculateDuration("Randomizing seed"):
135
+ if randomize_seed:
136
+ seed = random.randint(0, MAX_SEED)
137
+
138
+ image = generate_image(prompt, trigger_word, steps, seed, cfg_scale, width, height, lora_scale, progress)
139
+ pipe.to("cpu")
140
+ pipe.unload_lora_weights()
141
+ return image, seed
142
+
143
+ run_lora.zerogpu = True
144
+
145
+ css = '''
146
+ #gen_btn{height: 100%}
147
+ #title{text-align: center}
148
+ #title h1{font-size: 3em; display:inline-flex; align-items:center}
149
+ #title img{width: 100px; margin-right: 0.5em}
150
+ #gallery .grid-wrap{height: 10vh}
151
+ '''
152
+ with gr.Blocks(theme=gr.themes.Soft(), css=css) as app:
153
+ title = gr.HTML(
154
+ """<h1><img src="https://huggingface.co/AlekseyCalvin/HSTklimbimOPENfluxLora/resolve/main/acs62iv.png" alt="LoRA">OpenFlux LoRAsoon®</h1>""",
155
+ elem_id="title",
156
+ )
157
+ # Info blob stating what the app is running
158
+ info_blob = gr.HTML(
159
+ """<div id="info_blob"> SOON®'s curated LoRa Gallery & Art Manufactory Space.|Runs on Ostris' OpenFLUX.1 model + fast-gen LoRA & Zer0int's fine-tuned CLIP-GmP-ViT-L-14*! (*'normal' 77 tokens)| Largely stocked w/our trained LoRAs: Historic Color, Silver Age Poets, Sots Art, more!|</div>"""
160
+ )
161
+ # Info blob stating what the app is running
162
+ info_blob = gr.HTML(
163
+ """<div id="info_blob"> *Auto-planting of prompts with a choice LoRA trigger errors out in this space over flaws yet unclear. In its stead, we pose numbered LoRA-box rows & a matched token cheat-sheet: ungainly & free. So, prephrase your prompts w/: 1-2. HST style autochrome |3. RCA style Communist poster |4. SOTS art |5. HST Austin Osman Spare style |6. Vladimir Mayakovsky |7-8. Marina Tsvetaeva Tsvetaeva_02.CR2 |9. Anna Akhmatova |10. Osip Mandelshtam |11-12. Alexander Blok |13. Blok_02.CR2 |14. LEN Lenin |15. Leon Trotsky |16. Rosa Fluxemburg |17. HST Peterhof photo |18-19. HST |20. HST portrait |21. HST |22. HST 80s Perestroika-era Soviet photo |23-30. HST |31. How2Draw a__ |32. propaganda poster |33. TOK hybrid photo of__ with cartoon of__ |34. 2004 IMG_1099.CR2 photo |35. unexpected photo of |36. flmft |37. 80s yearbook photo |38. TOK portra |39. pficonics |40. retrofuturism |41. wh3r3sw4ld0 |42. amateur photo |43. crisp |44-45. IMG_1099.CR2 |46. FilmFotos |47. ff-collage |48. HST |49-50. AOS |51. cover </div>"""
164
+ )
165
+ selected_index = gr.State(None)
166
+ with gr.Row():
167
+ with gr.Column(scale=3):
168
+ prompt = gr.Textbox(label="Prompt", lines=1, placeholder="Select LoRa/Style & type prompt!")
169
+ with gr.Column(scale=1, elem_id="gen_column"):
170
+ generate_button = gr.Button("Generate", variant="primary", elem_id="gen_btn")
171
+ with gr.Row():
172
+ with gr.Column(scale=3):
173
+ selected_info = gr.Markdown("")
174
+ gallery = gr.Gallery(
175
+ [(item["image"], item["title"]) for item in loras],
176
+ label="LoRA Inventory",
177
+ allow_preview=False,
178
+ columns=3,
179
+ elem_id="gallery"
180
+ )
181
+
182
+ with gr.Column(scale=4):
183
+ result = gr.Image(label="Generated Image")
184
+
185
+ with gr.Row():
186
+ with gr.Accordion("Advanced Settings", open=True):
187
+ with gr.Column():
188
+ with gr.Row():
189
+ cfg_scale = gr.Slider(label="CFG Scale", minimum=1, maximum=20, step=1, value=3)
190
+ steps = gr.Slider(label="Steps", minimum=1, maximum=50, step=1, value=6)
191
+
192
+ with gr.Row():
193
+ width = gr.Slider(label="Width", minimum=256, maximum=1536, step=64, value=768)
194
+ height = gr.Slider(label="Height", minimum=256, maximum=1536, step=64, value=768)
195
+
196
+ with gr.Row():
197
+ randomize_seed = gr.Checkbox(True, label="Randomize seed")
198
+ seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0, randomize=True)
199
+ lora_scale = gr.Slider(label="LoRA Scale", minimum=0, maximum=1, step=0.01, value=0.95)
200
+
201
+ gallery.select(
202
+ update_selection,
203
+ inputs=[width, height],
204
+ outputs=[prompt, selected_info, selected_index, width, height]
205
+ )
206
+
207
+ gr.on(
208
+ triggers=[generate_button.click, prompt.submit],
209
+ fn=run_lora,
210
+ inputs=[prompt, cfg_scale, steps, selected_index, randomize_seed, seed, width, height, lora_scale],
211
+ outputs=[result, seed]
212
+ )
213
+
214
+ warnings.filterwarnings("ignore", category=FutureWarning)
215
+ app.queue(default_concurrency_limit=2).launch(show_error=True)
216
+ app.launch()