1inkusFace commited on
Commit
ff9a9d1
·
verified ·
1 Parent(s): 9a6c96c

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +439 -0
app.py ADDED
@@ -0,0 +1,439 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import subprocess
2
+ subprocess.run(['sh', './spaces.sh'])
3
+
4
+ import os
5
+
6
+ os.putenv('PYTORCH_NVML_BASED_CUDA_CHECK','1')
7
+ os.putenv('TORCH_LINALG_PREFER_CUSOLVER','1')
8
+ alloc_conf_parts = [
9
+ 'expandable_segments:True',
10
+ 'pinned_use_background_threads:True' # Specific to pinned memory.
11
+ ]
12
+ os.environ['PYTORCH_CUDA_ALLOC_CONF'] = ','.join(alloc_conf_parts)
13
+ os.environ["SAFETENSORS_FAST_GPU"] = "1"
14
+ os.putenv('HF_HUB_ENABLE_HF_TRANSFER','1')
15
+
16
+ import spaces
17
+
18
+ import gradio as gr
19
+ import numpy as np
20
+ import random
21
+
22
+
23
+
24
+ import torch
25
+ torch.backends.cuda.matmul.allow_tf32 = False
26
+ torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False
27
+ torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False
28
+ torch.backends.cudnn.allow_tf32 = False
29
+ torch.backends.cudnn.deterministic = False
30
+ torch.backends.cudnn.benchmark = False
31
+ torch.backends.cuda.preferred_blas_library="cublas"
32
+ torch.backends.cuda.preferred_linalg_library="cusolver"
33
+ torch.set_float32_matmul_precision("highest")
34
+
35
+ from diffusers import StableDiffusion3Pipeline, SD3Transformer2DModel, AutoencoderKL
36
+ from transformers import CLIPTextModelWithProjection, T5EncoderModel
37
+ from transformers import CLIPTokenizer, T5TokenizerFast
38
+
39
+ import re
40
+ import paramiko
41
+ import urllib
42
+ import time
43
+ from image_gen_aux import UpscaleWithModel
44
+ from huggingface_hub import hf_hub_download
45
+ import datetime
46
+ import cyper
47
+ from PIL import Image
48
+
49
+ #from accelerate import Accelerator
50
+
51
+ #accelerator = Accelerator(mixed_precision="bf16")
52
+
53
+ hftoken = os.getenv("HF_AUTH_TOKEN")
54
+
55
+ code = r'''
56
+ import torch
57
+ import paramiko
58
+ import os
59
+ import socket
60
+ import threading # NEW IMPORT
61
+ import queue # NEW IMPORT
62
+ FTP_HOST = 'noahcohn.com'
63
+ FTP_USER = 'ford442'
64
+ FTP_PASS = os.getenv("FTP_PASS")
65
+ FTP_DIR = 'img.noahcohn.com/stablediff/'
66
+ FTP_HOST_FALLBACK = '1ink.us'
67
+ FTP_DIR_FALLBACK = 'img.1ink.us/stablediff/'
68
+
69
+ # --- WORKER FUNCTION FOR THREADING ---
70
+ # This function contains the logic to connect to a single host.
71
+ # It will be executed by each of our threads.
72
+ def connect_worker(host, result_queue):
73
+ """Tries to connect to a single host and puts the successful transport object into the queue."""
74
+ transport = None
75
+ try:
76
+ transport = paramiko.Transport((host, 22))
77
+ # We still use the 5-second timeout for the handshake
78
+ transport.start_client(timeout=5)
79
+ transport.auth_password(username=FTP_USER, password=FTP_PASS)
80
+
81
+ # If we reach here, the connection was successful.
82
+ # Put the result in the queue for the main thread to use.
83
+ print(f"✅ Connection to {host} succeeded first.")
84
+ result_queue.put(transport)
85
+ except (paramiko.SSHException, socket.timeout, EOFError) as e:
86
+ # This is an expected failure, just print a note.
87
+ print(f"ℹ️ Connection to {host} failed or was too slow: {e}")
88
+ if transport:
89
+ transport.close()
90
+ except Exception as e:
91
+ # Handle any other unexpected errors.
92
+ print(f"❌ Unexpected error connecting to {host}: {e}")
93
+ if transport:
94
+ transport.close()
95
+
96
+ def upload_to_ftp(filename):
97
+ """
98
+ Attempts to connect to two FTP hosts simultaneously and uses the first one that responds.
99
+ It now uses a corresponding directory for the primary and fallback hosts.
100
+ """
101
+ hosts = [FTP_HOST]
102
+ if FTP_HOST_FALLBACK:
103
+ hosts.append(FTP_HOST_FALLBACK)
104
+
105
+ result_queue = queue.Queue()
106
+ threads = []
107
+
108
+ print(f"--> Racing connections to {hosts} for uploading {filename}...")
109
+ for host in hosts:
110
+ thread = threading.Thread(target=connect_worker, args=(host, result_queue))
111
+ thread.daemon = True
112
+ thread.start()
113
+ threads.append(thread)
114
+ try:
115
+ winning_transport = result_queue.get(timeout=7)
116
+
117
+ # --- THIS IS THE NEW LOGIC ---
118
+ # 1. Determine which host won the race.
119
+ winning_host = winning_transport.getpeername()[0]
120
+
121
+ # 2. Select the correct destination directory based on the winning host.
122
+ # If the fallback directory isn't specified, it safely defaults to the primary directory.
123
+ if winning_host == FTP_HOST:
124
+ destination_directory = FTP_DIR
125
+ else:
126
+ destination_directory = FTP_DIR_FALLBACK if FTP_DIR_FALLBACK else FTP_DIR
127
+ print(f"--> Proceeding with upload to {winning_host} in directory {destination_directory}...")
128
+ # 3. Construct the full destination path using the selected directory.
129
+ sftp = paramiko.SFTPClient.from_transport(winning_transport)
130
+ destination_path = os.path.join(destination_directory, os.path.basename(filename))
131
+ sftp.put(filename, destination_path)
132
+
133
+ print(f"✅ Successfully uploaded {filename}.")
134
+
135
+ sftp.close()
136
+ winning_transport.close()
137
+ except queue.Empty:
138
+ print("❌ Critical Error: Neither FTP host responded in time.")
139
+ except Exception as e:
140
+ print(f"❌ An unexpected error occurred during SFTP operation: {e}")
141
+ '''
142
+
143
+ pyx = cyper.inline(code, fast_indexing=True, directives=dict(boundscheck=False, wraparound=False, language_level=3))
144
+
145
+ device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
146
+
147
+ #vae=AutoencoderKL.from_pretrained("ford442/sdxl-vae-bf16", use_safetensors=True, subfolder='vae',token=True)
148
+ vaeX=AutoencoderKL.from_pretrained("ford442/stable-diffusion-3.5-large-fp32", safety_checker=None, use_safetensors=True, subfolder='vae', low_cpu_mem_usage=False, torch_dtype=torch.float32, token=True)
149
+ pipe = StableDiffusion3Pipeline.from_pretrained(
150
+ #"stabilityai # stable-diffusion-3.5-large",
151
+ "ford442/stable-diffusion-3.5-large-bf16",
152
+ trust_remote_code=True,
153
+ #vae=None,
154
+ #vae=AutoencoderKL.from_pretrained("ford442/stable-diffusion-3.5-large-fp32", use_safetensors=True, subfolder='vae',token=True),
155
+ #scheduler = FlowMatchHeunDiscreteScheduler.from_pretrained('ford442/stable-diffusion-3.5-large-bf16', subfolder='scheduler',token=True),
156
+ #text_encoder=None, #CLIPTextModelWithProjection.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder', token=True),
157
+ # text_encoder=CLIPTextModelWithProjection.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder', token=True),
158
+ #text_encoder_2=None, #CLIPTextModelWithProjection.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder_2',token=True),
159
+ # text_encoder_2=CLIPTextModelWithProjection.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder_2',token=True),
160
+ #text_encoder_3=None, #T5EncoderModel.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder_3',token=True),
161
+ # text_encoder_3=T5EncoderModel.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder_3',token=True),
162
+ #tokenizer=CLIPTokenizer.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", add_prefix_space=True, subfolder="tokenizer", token=True),
163
+ #tokenizer_2=CLIPTokenizer.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", add_prefix_space=True, subfolder="tokenizer_2", token=True),
164
+ transformer=None,
165
+ #tokenizer_3=T5TokenizerFast.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", add_prefix_space=False, use_fast=True, subfolder="tokenizer_3", token=True),
166
+ #torch_dtype=torch.bfloat16,
167
+ use_safetensors=True,
168
+ )
169
+ #text_encoder=CLIPTextModelWithProjection.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder', token=True).to(torch.device("cuda:0"), dtype=torch.bfloat16)
170
+ #text_encoder_2=CLIPTextModelWithProjection.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder_2',token=True).to(torch.device("cuda:0"), dtype=torch.bfloat16)
171
+ #text_encoder_3=T5EncoderModel.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder_3',token=True).to(torch.device("cuda:0"), dtype=torch.bfloat16)
172
+ ll_transformer=SD3Transformer2DModel.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='transformer',token=True).to(torch.device("cuda:0"), dtype=torch.bfloat16)
173
+ pipe.transformer=ll_transformer
174
+ pipe.load_lora_weights("ford442/sdxl-vae-bf16", weight_name="LoRA/UltraReal.safetensors")
175
+
176
+ #pipe.to(accelerator.device)
177
+ pipe.to(device=device, dtype=torch.bfloat16)
178
+
179
+ upscaler_2 = UpscaleWithModel.from_pretrained("Kim2091/ClearRealityV1").to(torch.device('cuda'))
180
+
181
+ MAX_SEED = np.iinfo(np.int32).max
182
+
183
+ MAX_IMAGE_SIZE = 4096
184
+
185
+ @spaces.GPU(duration=70)
186
+ def infer_60(
187
+ prompt,
188
+ negative_prompt_1,
189
+ negative_prompt_2,
190
+ negative_prompt_3,
191
+ width,
192
+ height,
193
+ guidance_scale,
194
+ num_inference_steps,
195
+ progress=gr.Progress(track_tqdm=True),
196
+ ):
197
+ seed = random.randint(0, MAX_SEED)
198
+ generator = torch.Generator(device='cuda').manual_seed(seed)
199
+ print('-- generating image --')
200
+ sd_image = pipe(
201
+ prompt=prompt,
202
+ prompt_2=prompt,
203
+ prompt_3=prompt,
204
+ negative_prompt=negative_prompt_1,
205
+ negative_prompt_2=negative_prompt_2,
206
+ negative_prompt_3=negative_prompt_3,
207
+ guidance_scale=guidance_scale,
208
+ num_inference_steps=num_inference_steps,
209
+ width=width,
210
+ height=height,
211
+ generator=generator,
212
+ max_sequence_length=512
213
+ ).images[0]
214
+ print('-- got image --')
215
+ timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
216
+ sd35_path = f"sd35ll_{timestamp}.png"
217
+ sd_image.save(sd35_path,optimize=False,compress_level=0)
218
+ pyx.upload_to_ftp(sd35_path)
219
+ with torch.no_grad():
220
+ upscale = upscaler_2(sd_image, tiling=True, tile_width=256, tile_height=256)
221
+ upscale2 = upscaler_2(upscale, tiling=True, tile_width=256, tile_height=256)
222
+ print('-- got upscaled image --')
223
+ downscale = upscale2.resize((upscale2.width // 4, upscale2.height // 4),Image.LANCZOS)
224
+ upscale_path = f"sd35ll_upscale_{timestamp}.png"
225
+ downscale.save(upscale_path,optimize=False,compress_level=0)
226
+ pyx.upload_to_ftp(upscale_path)
227
+ return sd_image, prompt
228
+
229
+ @spaces.GPU(duration=100)
230
+ def infer_90(
231
+ prompt,
232
+ negative_prompt_1,
233
+ negative_prompt_2,
234
+ negative_prompt_3,
235
+ width,
236
+ height,
237
+ guidance_scale,
238
+ num_inference_steps,
239
+ progress=gr.Progress(track_tqdm=True),
240
+ ):
241
+ seed = random.randint(0, MAX_SEED)
242
+ generator = torch.Generator(device='cuda').manual_seed(seed)
243
+ print('-- generating image --')
244
+ sd_image = pipe(
245
+ prompt=prompt,
246
+ prompt_2=prompt,
247
+ prompt_3=prompt,
248
+ negative_prompt=negative_prompt_1,
249
+ negative_prompt_2=negative_prompt_2,
250
+ negative_prompt_3=negative_prompt_3,
251
+ guidance_scale=guidance_scale,
252
+ num_inference_steps=num_inference_steps,
253
+ width=width,
254
+ height=height,
255
+ generator=generator,
256
+ max_sequence_length=512
257
+ ).images[0]
258
+ print('-- got image --')
259
+ timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
260
+ sd35_path = f"sd35ll_{timestamp}.png"
261
+ sd_image.save(sd35_path,optimize=False,compress_level=0)
262
+ pyx.upload_to_ftp(sd35_path)
263
+ with torch.no_grad():
264
+ upscale = upscaler_2(sd_image, tiling=True, tile_width=256, tile_height=256)
265
+ upscale2 = upscaler_2(upscale, tiling=True, tile_width=256, tile_height=256)
266
+ print('-- got upscaled image --')
267
+ downscale = upscale2.resize((upscale2.width // 4, upscale2.height // 4),Image.LANCZOS)
268
+ upscale_path = f"sd35ll_upscale_{timestamp}.png"
269
+ downscale.save(upscale_path,optimize=False,compress_level=0)
270
+ pyx.upload_to_ftp(upscale_path)
271
+ return sd_image, prompt
272
+
273
+ @spaces.GPU(duration=120)
274
+ def infer_110(
275
+ prompt,
276
+ negative_prompt_1,
277
+ negative_prompt_2,
278
+ negative_prompt_3,
279
+ width,
280
+ height,
281
+ guidance_scale,
282
+ num_inference_steps,
283
+ progress=gr.Progress(track_tqdm=True),
284
+ ):
285
+ seed = random.randint(0, MAX_SEED)
286
+ generator = torch.Generator(device='cuda').manual_seed(seed)
287
+ print('-- generating image --')
288
+ sd_image = pipe(
289
+ prompt=prompt,
290
+ prompt_2=prompt,
291
+ prompt_3=prompt,
292
+ negative_prompt=negative_prompt_1,
293
+ negative_prompt_2=negative_prompt_2,
294
+ negative_prompt_3=negative_prompt_3,
295
+ guidance_scale=guidance_scale,
296
+ num_inference_steps=num_inference_steps,
297
+ width=width,
298
+ height=height,
299
+ generator=generator,
300
+ max_sequence_length=512
301
+ ).images[0]
302
+ print('-- got image --')
303
+ timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
304
+ sd35_path = f"sd35ll_{timestamp}.png"
305
+ sd_image.save(sd35_path,optimize=False,compress_level=0)
306
+ pyx.upload_to_ftp(sd35_path)
307
+ with torch.no_grad():
308
+ upscale = upscaler_2(sd_image, tiling=True, tile_width=256, tile_height=256)
309
+ upscale2 = upscaler_2(upscale, tiling=True, tile_width=256, tile_height=256)
310
+ print('-- got upscaled image --')
311
+ downscale = upscale2.resize((upscale2.width // 4, upscale2.height // 4),Image.LANCZOS)
312
+ upscale_path = f"sd35ll_upscale_{timestamp}.png"
313
+ downscale.save(upscale_path,optimize=False,compress_level=0)
314
+ pyx.upload_to_ftp(upscale_path)
315
+ return sd_image, prompt
316
+
317
+ css = """
318
+ #col-container {margin: 0 auto;max-width: 640px;}
319
+ body{background-color: blue;}
320
+ """
321
+
322
+ with gr.Blocks(theme=gr.themes.Origin(),css=css) as demo:
323
+ with gr.Column(elem_id="col-container"):
324
+ gr.Markdown(" # StableDiffusion 3.5 Large with UltraReal lora test")
325
+ expanded_prompt_output = gr.Textbox(label="Prompt", lines=1) # Add this line
326
+ with gr.Row():
327
+ prompt = gr.Text(
328
+ label="Prompt",
329
+ show_label=False,
330
+ max_lines=1,
331
+ placeholder="Enter your prompt",
332
+ container=False,
333
+ )
334
+ run_button_60 = gr.Button("Run 60", scale=0, variant="primary")
335
+ run_button_90 = gr.Button("Run 90", scale=0, variant="primary")
336
+ run_button_110 = gr.Button("Run 110", scale=0, variant="primary")
337
+ result = gr.Image(label="Result", show_label=False)
338
+ with gr.Accordion("Advanced Settings", open=True):
339
+ negative_prompt_1 = gr.Text(
340
+ label="Negative prompt 1",
341
+ max_lines=1,
342
+ placeholder="Enter a negative prompt",
343
+ visible=True,
344
+ value="bad anatomy, poorly drawn hands, distorted face, blurry, out of frame, low resolution, grainy, pixelated, disfigured, mutated, extra limbs, bad composition"
345
+ )
346
+ negative_prompt_2 = gr.Text(
347
+ label="Negative prompt 2",
348
+ max_lines=1,
349
+ placeholder="Enter a second negative prompt",
350
+ visible=True,
351
+ value="unrealistic, cartoon, anime, sketch, painting, drawing, illustration, graphic, digital art, render, 3d, blurry, deformed, disfigured, poorly drawn, bad anatomy, mutated, extra limbs, ugly, out of frame, bad composition, low resolution, grainy, pixelated, noisy, oversaturated, undersaturated, (worst quality, low quality:1.3), (bad hands, missing fingers:1.2)"
352
+ )
353
+ negative_prompt_3 = gr.Text(
354
+ label="Negative prompt 3",
355
+ max_lines=1,
356
+ placeholder="Enter a third negative prompt",
357
+ visible=True,
358
+ value="(worst quality, low quality:1.3), (bad anatomy, bad hands, missing fingers, extra digit, fewer digits:1.2), (blurry:1.1), cropped, watermark, text, signature, logo, jpeg artifacts, (ugly, deformed, disfigured:1.2), (poorly drawn:1.2), mutated, extra limbs, (bad proportions, gross proportions:1.2), (malformed limbs, missing arms, missing legs, extra arms, extra legs:1.2), (fused fingers, too many fingers, long neck:1.2), (unnatural body, unnatural pose:1.1), out of frame, (bad composition, poorly composed:1.1), (oversaturated, undersaturated:1.1), (grainy, pixelated:1.1), (low resolution, noisy:1.1), (unrealistic, distorted:1.1), (extra fingers, mutated hands, poorly drawn hands, bad hands:1.3), (missing fingers:1.3)"
359
+ )
360
+ num_iterations = gr.Number(
361
+ value=1000,
362
+ label="Number of Iterations")
363
+ with gr.Row():
364
+ width = gr.Slider(
365
+ label="Width",
366
+ minimum=256,
367
+ maximum=MAX_IMAGE_SIZE,
368
+ step=32,
369
+ value=768,
370
+ )
371
+ height = gr.Slider(
372
+ label="Height",
373
+ minimum=256,
374
+ maximum=MAX_IMAGE_SIZE,
375
+ step=32,
376
+ value=768,
377
+ )
378
+ guidance_scale = gr.Slider(
379
+ label="Guidance scale",
380
+ minimum=0.0,
381
+ maximum=30.0,
382
+ step=0.1,
383
+ value=4.2,
384
+ )
385
+ num_inference_steps = gr.Slider(
386
+ label="Number of inference steps",
387
+ minimum=1,
388
+ maximum=500,
389
+ step=1,
390
+ value=100,
391
+ )
392
+ gr.on(
393
+ triggers=[run_button_60.click, prompt.submit],
394
+ fn=infer_60,
395
+ inputs=[
396
+ prompt,
397
+ negative_prompt_1,
398
+ negative_prompt_2,
399
+ negative_prompt_3,
400
+ width,
401
+ height,
402
+ guidance_scale,
403
+ num_inference_steps,
404
+ ],
405
+ outputs=[result, expanded_prompt_output],
406
+ )
407
+ gr.on(
408
+ triggers=[run_button_90.click, prompt.submit],
409
+ fn=infer_90,
410
+ inputs=[
411
+ prompt,
412
+ negative_prompt_1,
413
+ negative_prompt_2,
414
+ negative_prompt_3,
415
+ width,
416
+ height,
417
+ guidance_scale,
418
+ num_inference_steps,
419
+ ],
420
+ outputs=[result, expanded_prompt_output],
421
+ )
422
+ gr.on(
423
+ triggers=[run_button_110.click, prompt.submit],
424
+ fn=infer_110,
425
+ inputs=[
426
+ prompt,
427
+ negative_prompt_1,
428
+ negative_prompt_2,
429
+ negative_prompt_3,
430
+ width,
431
+ height,
432
+ guidance_scale,
433
+ num_inference_steps,
434
+ ],
435
+ outputs=[result, expanded_prompt_output],
436
+ )
437
+
438
+ if __name__ == "__main__":
439
+ demo.launch()