add files

img/dev-requirements.txt

@@ -0,0 +1,11 @@
+pytest
+
+pytest-asyncio
+requests-futures==1.0.0
+httpx
+djlint
+pytest-env==0.8.1
+ipython
+
+line-profiler-pycharm==1.1.0
+line-profiler==4.0.3

img/env.py

@@ -0,0 +1,2 @@
+BUCKET_NAME = 'static.netwrck.com'
+BUCKET_PATH = 'static/uploads'

img/img2img.py

@@ -0,0 +1,25 @@
+import requests
+import torch
+from PIL import Image
+from io import BytesIO
+
+from diffusers import StableDiffusionImg2ImgPipeline
+
+device = "cuda"
+model_id_or_path = "runwayml/stable-diffusion-v1-5"
+# model_id_or_path = "models/stable-diffusion-xl-base-0.9"
+pipe = StableDiffusionImg2ImgPipeline.from_pretrained(model_id_or_path, torch_dtype=torch.float16, variant="fp16", safety_checker=None)
+pipe = pipe.to(device)
+
+url = "https://raw.githubusercontent.com/CompVis/stable-diffusion/main/assets/stable-samples/img2img/sketch-mountains-input.jpg"
+
+response = requests.get(url)
+# init_image = Image.open(BytesIO(response.content)).convert("RGB")
+init_image = Image.open("/mnt/c/Users/leepenkman/Pictures/aiknight-neon-punk-fantasy-art-good-looking-trending-fantastic-1.webp").convert("RGB")
+# init_image = init_image.resize((768, 512))
+init_image = init_image.resize((1920, 1080))
+
+prompt = "knight neon punk fantasy art good looking trending fantastic"
+
+images = pipe(prompt=prompt, image=init_image, strength=0.75, guidance_scale=7.5).images
+images[0].save("fantasy_landscape.png")

img/img2imgsd.py

@@ -0,0 +1,74 @@
+from pathlib import Path
+
+import numpy as np
+import requests
+import torch
+from PIL import Image
+from io import BytesIO
+
+# from diffusers import StableDiffusionImg2ImgPipeline
+
+# device = "cuda"
+# model_id_or_path = "runwayml/stable-diffusion-v1-5"
+# # model_id_or_path = "models/stable-diffusion-xl-base-0.9"
+# pipe = StableDiffusionImg2ImgPipeline.from_pretrained(model_id_or_path, torch_dtype=torch.float16, variant="fp16", safety_checker=None)
+# pipe = pipe.to(device)
+
+from diffusers import StableDiffusionXLImg2ImgPipeline
+from diffusers.utils import load_image
+
+from stable_diffusion_server.utils import log_time
+
+pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained(
+"stabilityai/stable-diffusion-xl-refiner-1.0",
+# "models/stable-diffusion-xl-base-0.9",
+    torch_dtype = torch.float16,
+    use_safetensors=True,
+    variant="fp16",
+)
+pipe = pipe.to("cuda") #  # "LayerNormKernelImpl" not implemented for 'Half' error if its on cpu it cant do fp16
+# idea composite: and re prompt img-img to support different sizes
+
+# url = "https://raw.githubusercontent.com/CompVis/stable-diffusion/main/assets/stable-samples/img2img/sketch-mountains-input.jpg"
+#
+# response = requests.get(url)
+# init_image = Image.open(BytesIO(response.content)).convert("RGB")
+# init_image = init_image.resize((768, 512))
+# successfully inpaints a deleted area strength=0.75
+# init_image = Image.open("/mnt/c/Users/leepenkman/Pictures/aiart/ainostalgic-colorful-relaxing-chill-realistic-cartoon-Charcoal-illustration-fantasy-fauvist-abstract-impressionist-watercolor-painting-Background-location-scenery-amazing-wonderful-Dog-Shelter-Worker-Dog.webp").convert("RGB")
+# redo something? strength 1
+# init_image = Image.open("/home/lee/code/sdif/mask.png").convert("RGB")
+init_image = Image.open("/mnt/c/Users/leepenkman/Pictures/dogstretch.png").convert("RGB")
+# init_image = Image.open("/mnt/c/Users/leepenkman/Pictures/dogcenter.png").convert("RGB")
+
+# init_image = init_image.resize((1080, 1920))
+init_image = init_image.resize((1920, 1080))
+# init_image = init_image.resize((1024, 1024))
+
+prompt = "A fantasy landscape, trending on artstation, beautiful amazing unreal surreal gorgeous impressionism"
+prompt = "mouth open nostalgic colorful relaxing chill realistic cartoon Charcoal illustration fantasy fauvist abstract impressionist watercolor painting Background location scenery amazing wonderful Dog Shelter Worker Dog"
+
+# images = pipe(prompt=prompt, image=init_image, strength=0.75, guidance_scale=7.5).images
+# images[0].save("fantasy_landscape.png")
+#
+# # url = "https://huggingface.co/datasets/patrickvonplaten/images/resolve/main/aa_xl/000000009.png"
+#
+# init_image = load_image(url).convert("RGB")
+# prompt = "a photo of an astronaut riding a horse on mars"
+study_dir = "images/study2"
+Path(study_dir).mkdir(parents=True, exist_ok=True)
+
+with log_time("img2img"):
+    with torch.inference_mode():
+        # for strength in range(.1, 1, .1):
+        for strength in np.linspace(.1, 1, 10):
+            image = pipe(prompt=prompt, image=init_image, strength=strength, guidance_scale=7.6).images[0]
+            image.save(
+                study_dir + "/fantasy_dogimgimgdogstretchopening" + str(strength) + "guidance_scale" + str(7.6) + ".png")
+        #     # for guidance_scale in range(1, 10, .5):
+        #     for guidance_scale in np.linspace(1, 100, 10):
+        #         image = pipe(prompt=prompt, image=init_image, strength=strength, guidance_scale=guidance_scale).images[0]
+        #         image.save("images/study/fantasy_dogimgimgdogstretch" + str(strength) + "guidance_scale" + str(guidance_scale) + ".png")
+        # image = pipe(prompt, image=init_image, strength=0.2, guidance_scale=7.5).images[0]
+        # image.save("images/fantasy_dogimgimgdogstretch.png")
+        # image.save("images/fantasy_dogimgimgdogcenter.png")

img/img2imgsdr.py

@@ -0,0 +1,53 @@
+import PIL.Image
+
+from diffusers import DiffusionPipeline
+import torch
+
+import numpy as np
+
+from stable_diffusion_server.utils import log_time
+
+pipe = DiffusionPipeline.from_pretrained(
+    "models/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, variant="fp16", use_safetensors=True
+)
+pipe.to("cuda")
+
+refiner = DiffusionPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-refiner-1.0",
+    text_encoder_2=pipe.text_encoder_2,
+    vae=pipe.vae,
+    torch_dtype=torch.float16,
+    use_safetensors=True,
+    variant="fp16",
+)
+refiner.to("cuda")
+
+prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
+use_refiner = True
+with log_time('diffuse'):
+    with torch.inference_mode():
+        image = pipe(prompt=prompt, output_type="latent" if use_refiner else "pil").images[0]
+        # experiment try deleting a whole bunch of pixels and see if the refiner can recreate them
+        # delete top 30% of pixels
+        # image = image[0:0.7]
+        #pixels to delete
+        # pixels_to_delete = int(0.3 * 1024)
+        # delete top 30% of pixels
+        # image.save("latent.png")
+        # image_data = PIL.Image.fromarray(image)
+        # image_data.save("latent.png")
+
+        # image = np.array(image)
+        pixels_to_delete = int(0.3 * image.shape[0])
+        idx_to_delete = np.ones(image.shape[0], dtype=bool, device="cuda")
+        idx_to_delete[:pixels_to_delete] = False
+        image[idx_to_delete] = [0,0,0]
+
+        # image_data = PIL.Image.fromarray(image)
+        # image_data.save("latentcleared.png")
+
+
+        image = refiner(prompt=prompt, image=image[None, :]).images[0]
+
+
+

img/inpaint.py

@@ -0,0 +1,62 @@
+import torch
+
+from diffusers import StableDiffusionXLInpaintPipeline
+from diffusers.utils import load_image
+
+from stable_diffusion_server.utils import log_time
+
+import numpy as np
+import PIL.Image
+
+pipe = StableDiffusionXLInpaintPipeline.from_pretrained(
+    "models/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, variant="fp16", use_safetensors=True
+)
+pipe.to("cuda")
+
+refiner = StableDiffusionXLInpaintPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-refiner-1.0",
+    text_encoder_2=pipe.text_encoder_2,
+    vae=pipe.vae,
+    torch_dtype=torch.float16,
+    use_safetensors=True,
+    variant="fp16",
+)
+refiner.to("cuda")
+
+img_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo.png"
+mask_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo_mask.png"
+# inpaint_and_upload_image?prompt=majestic tiger sitting on a bench&image_url=https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo.png&mask_url=https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo_mask.png&save_path=tests/inpaint.webp
+# inpainting can be used to upscale to 1080p
+
+
+init_image = load_image(img_url).convert("RGB")
+# mask_image = load_image(mask_url).convert("RGB")
+# mask image all ones same shape as init_image
+
+# here's a failed experiment: inpainting cannot be used as style transfer/it doesnt recreate ain image doing a full mask in this way
+image_size = init_image.size
+ones_of_size = np.ones(image_size, np.uint8) * 255
+mask_image = PIL.Image.fromarray(ones_of_size.astype(np.uint8))
+# mask_image = torch.ones_like(init_image) * 255
+prompt = "A majestic tiger sitting on a bench, castle backdrop elegent anime"
+num_inference_steps = 75
+high_noise_frac = 0.7
+with log_time("inpaint"):
+    with torch.inference_mode():
+        image = pipe(
+            prompt=prompt,
+            image=init_image,
+            mask_image=mask_image,
+            num_inference_steps=num_inference_steps,
+            denoising_start=high_noise_frac,
+            output_type="latent",
+        ).images
+        image = refiner(
+            prompt=prompt,
+            image=image,
+            mask_image=mask_image,
+            num_inference_steps=num_inference_steps,
+            denoising_start=high_noise_frac,
+        ).images[0]
+
+image.save("inpaintfull.png")

img/log.0925

@@ -0,0 +1,53 @@
+v-haipe+   551 16041 99 08:16 pts/2    00:00:17 python LiLa/gsm8k_cluster.py
+v-haipe+  9211 10235  3 Sep24 pts/10   00:32:12 python LiLa/chatgpt_evol_lila_gsm8k_domain.py --start 0 --end 2000
+v-haipe+  9288 10459  3 Sep24 pts/11   00:28:30 python LiLa/chatgpt_evol_lila_gsm8k_domain.py --start 2000 --end 4000
+v-haipe+  9310 10667  3 Sep24 pts/12   00:27:45 python LiLa/chatgpt_evol_lila_gsm8k_domain.py --start 4000 --end 6000
+v-haipe+  9341 10865  3 Sep24 pts/13   00:26:50 python LiLa/chatgpt_evol_lila_gsm8k_domain.py --start 6000 --end 8000
+v-haipe+  9379 25248  3 Sep24 pts/16   00:27:01 python LiLa/chatgpt_evol_lila_gsm8k_domain.py --start 8000 --end 10000
+v-haipe+  9410 25467  3 Sep24 pts/17   00:27:17 python LiLa/chatgpt_evol_lila_gsm8k_domain.py --start 10000 --end 12000
+v-haipe+  9438 26561  3 Sep24 pts/19   00:27:17 python LiLa/chatgpt_evol_lila_gsm8k_domain.py --start 12000 --end 14000
+v-haipe+  9469 26761  3 Sep24 pts/20   00:26:55 python LiLa/chatgpt_evol_lila_gsm8k_domain.py --start 14000 --end 16000
+v-haipe+  9500 26968  3 Sep24 pts/21   00:27:09 python LiLa/chatgpt_evol_lila_gsm8k_domain.py --start 16000 --end 18000
+v-haipe+  9531 27172  3 Sep24 pts/22   00:29:29 python LiLa/chatgpt_evol_lila_gsm8k_domain.py --start 18000 --end 20000
+v-haipe+  9775  9560  3 Sep24 pts/29   00:30:29 python LiLa/chatgpt_evol_lila_gsm8k_domain.py --start 20000 --end 22000
+v-haipe+ 11262 24577  0 Sep23 pts/8    00:00:06 python app.py
+v-haipe+ 11300 11262  0 Sep23 pts/8    00:20:54 /home/v-haipengluo/.conda/envs/wizardweb/bin/python /workspaceblobstore/qins/test/20220316/kai/research/code_repo/wizard_verse/code_repo/server_code/wizard_verse/lm/server_lm/app.py
+v-haipe+ 11604 20782 98 Sep23 pts/4    2-00:06:57 python -m vllm.entrypoints.api_server --model /workspaceblobstore/caxu/trained_models/13Bv2_497kcontinueroleplay_dsys_2048_e4_2e_5/checkpoint-75 --host phlrr3006.guest.corp.microsoft.com --port 7991
+v-haipe+ 13722 22601  0 Sep24 pts/6    00:09:37 /home/v-haipengluo/.conda/envs/sdxl/bin/python /home/v-haipengluo/.conda/envs/sdxl/bin/uvicorn --host=phlrr3006.guest.corp.microsoft.com --port 7999 --workers 1 --backlog 1 --limit-concurrency 4 main_v3:app
+v-haipe+ 13830 13722  0 Sep24 pts/6    00:00:05 /home/v-haipengluo/.conda/envs/sdxl/bin/python /home/v-haipengluo/.conda/envs/sdxl/bin/uvicorn --host=phlrr3006.guest.corp.microsoft.com --port 7999 --workers 1 --backlog 1 --limit-concurrency 4 main_v3:app
+v-haipe+ 13834 13722  0 Sep24 pts/6    00:00:05 /home/v-haipengluo/.conda/envs/sdxl/bin/python /home/v-haipengluo/.conda/envs/sdxl/bin/uvicorn --host=phlrr3006.guest.corp.microsoft.com --port 7999 --workers 1 --backlog 1 --limit-concurrency 4 main_v3:app
+v-haipe+ 13837 13722  0 Sep24 pts/6    00:00:05 /home/v-haipengluo/.conda/envs/sdxl/bin/python /home/v-haipengluo/.conda/envs/sdxl/bin/uvicorn --host=phlrr3006.guest.corp.microsoft.com --port 7999 --workers 1 --backlog 1 --limit-concurrency 4 main_v3:app
+v-haipe+ 13839 13722  0 Sep24 pts/6    00:00:05 /home/v-haipengluo/.conda/envs/sdxl/bin/python /home/v-haipengluo/.conda/envs/sdxl/bin/uvicorn --host=phlrr3006.guest.corp.microsoft.com --port 7999 --workers 1 --backlog 1 --limit-concurrency 4 main_v3:app
+v-haipe+ 13841 13722  0 Sep24 pts/6    00:00:05 /home/v-haipengluo/.conda/envs/sdxl/bin/python /home/v-haipengluo/.conda/envs/sdxl/bin/uvicorn --host=phlrr3006.guest.corp.microsoft.com --port 7999 --workers 1 --backlog 1 --limit-concurrency 4 main_v3:app
+v-haipe+ 13843 13722  0 Sep24 pts/6    00:00:05 /home/v-haipengluo/.conda/envs/sdxl/bin/python /home/v-haipengluo/.conda/envs/sdxl/bin/uvicorn --host=phlrr3006.guest.corp.microsoft.com --port 7999 --workers 1 --backlog 1 --limit-concurrency 4 main_v3:app
+v-haipe+ 13845 13722  0 Sep24 pts/6    00:00:05 /home/v-haipengluo/.conda/envs/sdxl/bin/python /home/v-haipengluo/.conda/envs/sdxl/bin/uvicorn --host=phlrr3006.guest.corp.microsoft.com --port 7999 --workers 1 --backlog 1 --limit-concurrency 4 main_v3:app
+v-haipe+ 13847 13722  0 Sep24 pts/6    00:00:05 /home/v-haipengluo/.conda/envs/sdxl/bin/python /home/v-haipengluo/.conda/envs/sdxl/bin/uvicorn --host=phlrr3006.guest.corp.microsoft.com --port 7999 --workers 1 --backlog 1 --limit-concurrency 4 main_v3:app
+v-haipe+ 13849 13722  0 Sep24 pts/6    00:00:05 /home/v-haipengluo/.conda/envs/sdxl/bin/python /home/v-haipengluo/.conda/envs/sdxl/bin/uvicorn --host=phlrr3006.guest.corp.microsoft.com --port 7999 --workers 1 --backlog 1 --limit-concurrency 4 main_v3:app
+v-haipe+ 13851 13722  0 Sep24 pts/6    00:00:05 /home/v-haipengluo/.conda/envs/sdxl/bin/python /home/v-haipengluo/.conda/envs/sdxl/bin/uvicorn --host=phlrr3006.guest.corp.microsoft.com --port 7999 --workers 1 --backlog 1 --limit-concurrency 4 main_v3:app
+v-haipe+ 13853 13722  0 Sep24 pts/6    00:00:05 /home/v-haipengluo/.conda/envs/sdxl/bin/python /home/v-haipengluo/.conda/envs/sdxl/bin/uvicorn --host=phlrr3006.guest.corp.microsoft.com --port 7999 --workers 1 --backlog 1 --limit-concurrency 4 main_v3:app
+v-haipe+ 13855 13722  0 Sep24 pts/6    00:00:05 /home/v-haipengluo/.conda/envs/sdxl/bin/python /home/v-haipengluo/.conda/envs/sdxl/bin/uvicorn --host=phlrr3006.guest.corp.microsoft.com --port 7999 --workers 1 --backlog 1 --limit-concurrency 4 main_v3:app
+v-haipe+ 13857 13722  0 Sep24 pts/6    00:00:05 /home/v-haipengluo/.conda/envs/sdxl/bin/python /home/v-haipengluo/.conda/envs/sdxl/bin/uvicorn --host=phlrr3006.guest.corp.microsoft.com --port 7999 --workers 1 --backlog 1 --limit-concurrency 4 main_v3:app
+v-haipe+ 13859 13722  0 Sep24 pts/6    00:00:05 /home/v-haipengluo/.conda/envs/sdxl/bin/python /home/v-haipengluo/.conda/envs/sdxl/bin/uvicorn --host=phlrr3006.guest.corp.microsoft.com --port 7999 --workers 1 --backlog 1 --limit-concurrency 4 main_v3:app
+v-haipe+ 13861 13722  0 Sep24 pts/6    00:00:05 /home/v-haipengluo/.conda/envs/sdxl/bin/python /home/v-haipengluo/.conda/envs/sdxl/bin/uvicorn --host=phlrr3006.guest.corp.microsoft.com --port 7999 --workers 1 --backlog 1 --limit-concurrency 4 main_v3:app
+v-haipe+ 13863 13722  0 Sep24 pts/6    00:00:05 /home/v-haipengluo/.conda/envs/sdxl/bin/python /home/v-haipengluo/.conda/envs/sdxl/bin/uvicorn --host=phlrr3006.guest.corp.microsoft.com --port 7999 --workers 1 --backlog 1 --limit-concurrency 4 main_v3:app
+v-haipe+ 13865 13722  0 Sep24 pts/6    00:00:05 /home/v-haipengluo/.conda/envs/sdxl/bin/python /home/v-haipengluo/.conda/envs/sdxl/bin/uvicorn --host=phlrr3006.guest.corp.microsoft.com --port 7999 --workers 1 --backlog 1 --limit-concurrency 4 main_v3:app
+v-haipe+ 13867 13722  0 Sep24 pts/6    00:00:05 /home/v-haipengluo/.conda/envs/sdxl/bin/python /home/v-haipengluo/.conda/envs/sdxl/bin/uvicorn --host=phlrr3006.guest.corp.microsoft.com --port 7999 --workers 1 --backlog 1 --limit-concurrency 4 main_v3:app
+v-haipe+ 13869 13722  0 Sep24 pts/6    00:00:05 /home/v-haipengluo/.conda/envs/sdxl/bin/python /home/v-haipengluo/.conda/envs/sdxl/bin/uvicorn --host=phlrr3006.guest.corp.microsoft.com --port 7999 --workers 1 --backlog 1 --limit-concurrency 4 main_v3:app
+v-haipe+ 13871 13722  0 Sep24 pts/6    00:00:05 /home/v-haipengluo/.conda/envs/sdxl/bin/python /home/v-haipengluo/.conda/envs/sdxl/bin/uvicorn --host=phlrr3006.guest.corp.microsoft.com --port 7999 --workers 1 --backlog 1 --limit-concurrency 4 main_v3:app
+v-haipe+ 13873 13722  0 Sep24 pts/6    00:00:05 /home/v-haipengluo/.conda/envs/sdxl/bin/python /home/v-haipengluo/.conda/envs/sdxl/bin/uvicorn --host=phlrr3006.guest.corp.microsoft.com --port 7999 --workers 1 --backlog 1 --limit-concurrency 4 main_v3:app
+v-haipe+ 13875 13722  0 Sep24 pts/6    00:00:05 /home/v-haipengluo/.conda/envs/sdxl/bin/python /home/v-haipengluo/.conda/envs/sdxl/bin/uvicorn --host=phlrr3006.guest.corp.microsoft.com --port 7999 --workers 1 --backlog 1 --limit-concurrency 4 main_v3:app
+v-haipe+ 13877 13722  0 Sep24 pts/6    00:00:05 /home/v-haipengluo/.conda/envs/sdxl/bin/python /home/v-haipengluo/.conda/envs/sdxl/bin/uvicorn --host=phlrr3006.guest.corp.microsoft.com --port 7999 --workers 1 --backlog 1 --limit-concurrency 4 main_v3:app
+v-haipe+ 13879 13722  0 Sep24 pts/6    00:00:05 /home/v-haipengluo/.conda/envs/sdxl/bin/python /home/v-haipengluo/.conda/envs/sdxl/bin/uvicorn --host=phlrr3006.guest.corp.microsoft.com --port 7999 --workers 1 --backlog 1 --limit-concurrency 4 main_v3:app
+v-haipe+ 13881 13722  0 Sep24 pts/6    00:00:05 /home/v-haipengluo/.conda/envs/sdxl/bin/python /home/v-haipengluo/.conda/envs/sdxl/bin/uvicorn --host=phlrr3006.guest.corp.microsoft.com --port 7999 --workers 1 --backlog 1 --limit-concurrency 4 main_v3:app
+v-haipe+ 13883 13722  0 Sep24 pts/6    00:00:05 /home/v-haipengluo/.conda/envs/sdxl/bin/python /home/v-haipengluo/.conda/envs/sdxl/bin/uvicorn --host=phlrr3006.guest.corp.microsoft.com --port 7999 --workers 1 --backlog 1 --limit-concurrency 4 main_v3:app
+v-haipe+ 13885 13722  0 Sep24 pts/6    00:00:05 /home/v-haipengluo/.conda/envs/sdxl/bin/python /home/v-haipengluo/.conda/envs/sdxl/bin/uvicorn --host=phlrr3006.guest.corp.microsoft.com --port 7999 --workers 1 --backlog 1 --limit-concurrency 4 main_v3:app
+v-haipe+ 13887 13722  0 Sep24 pts/6    00:00:05 /home/v-haipengluo/.conda/envs/sdxl/bin/python /home/v-haipengluo/.conda/envs/sdxl/bin/uvicorn --host=phlrr3006.guest.corp.microsoft.com --port 7999 --workers 1 --backlog 1 --limit-concurrency 4 main_v3:app
+v-haipe+ 18319 15852  0 05:34 pts/1    00:00:03 /home/v-haipengluo/.conda/envs/llamax/bin/python /home/v-haipengluo/.conda/envs/llamax/bin/deepspeed --master_port 29500 --hostfile=hostfile --include=localhost:1,3,4,5,6,7 src/train.py --model_name_or_path /workspaceblobstore/qins/test/20220316/haipeng/output_weights/llamax_13b_stackexchange_MATH_12w_sample_5w_score0.5_trainset_2e-5/checkpoint-992 --data_path /workspaceblobstore/qins/test/20220316/haipeng/data/Math_datasets/MATH_the_answer_is_format/hendrycks_math_7500_ori_gpt4_ori_15k.json --output_dir /workspaceblobstore/qins/test/20220316/haipeng/output_weights/llamax_13b_continue_train_stackMATH5w_checkpoint992_hendrycks_math_7500_ori_gpt4_ori_15k --num_train_epochs 3 --model_max_length 1150 --per_device_train_batch_size 17 --per_device_eval_batch_size 1 --gradient_accumulation_steps 1 --evaluation_strategy no --save_strategy steps --save_steps 36 --save_total_limit 200 --learning_rate 2e-5 --warmup_steps 10 --logging_steps 2 --lr_scheduler_type cosine --report_to tensorboard --gradient_checkpointing True --deepspeed src/configs/deepspeed_config.json --fp16 True
+v-haipe+ 18333 18319  0 05:34 pts/1    00:00:03 /home/v-haipengluo/.conda/envs/llamax/bin/python -u -m deepspeed.launcher.launch --world_info=eyJsb2NhbGhvc3QiOiBbMSwgMywgNCwgNSwgNiwgN119 --master_addr=127.0.0.1 --master_port=29500 --enable_each_rank_log=None src/train.py --model_name_or_path /workspaceblobstore/qins/test/20220316/haipeng/output_weights/llamax_13b_stackexchange_MATH_12w_sample_5w_score0.5_trainset_2e-5/checkpoint-992 --data_path /workspaceblobstore/qins/test/20220316/haipeng/data/Math_datasets/MATH_the_answer_is_format/hendrycks_math_7500_ori_gpt4_ori_15k.json --output_dir /workspaceblobstore/qins/test/20220316/haipeng/output_weights/llamax_13b_continue_train_stackMATH5w_checkpoint992_hendrycks_math_7500_ori_gpt4_ori_15k --num_train_epochs 3 --model_max_length 1150 --per_device_train_batch_size 17 --per_device_eval_batch_size 1 --gradient_accumulation_steps 1 --evaluation_strategy no --save_strategy steps --save_steps 36 --save_total_limit 200 --learning_rate 2e-5 --warmup_steps 10 --logging_steps 2 --lr_scheduler_type cosine --report_to tensorboard --gradient_checkpointing True --deepspeed src/configs/deepspeed_config.json --fp16 True
+v-haipe+ 18346 18333 99 05:34 pts/1    03:20:42 /home/v-haipengluo/.conda/envs/llamax/bin/python -u src/train.py --local_rank=0 --model_name_or_path /workspaceblobstore/qins/test/20220316/haipeng/output_weights/llamax_13b_stackexchange_MATH_12w_sample_5w_score0.5_trainset_2e-5/checkpoint-992 --data_path /workspaceblobstore/qins/test/20220316/haipeng/data/Math_datasets/MATH_the_answer_is_format/hendrycks_math_7500_ori_gpt4_ori_15k.json --output_dir /workspaceblobstore/qins/test/20220316/haipeng/output_weights/llamax_13b_continue_train_stackMATH5w_checkpoint992_hendrycks_math_7500_ori_gpt4_ori_15k --num_train_epochs 3 --model_max_length 1150 --per_device_train_batch_size 17 --per_device_eval_batch_size 1 --gradient_accumulation_steps 1 --evaluation_strategy no --save_strategy steps --save_steps 36 --save_total_limit 200 --learning_rate 2e-5 --warmup_steps 10 --logging_steps 2 --lr_scheduler_type cosine --report_to tensorboard --gradient_checkpointing True --deepspeed src/configs/deepspeed_config.json --fp16 True
+v-haipe+ 18347 18333 99 05:34 pts/1    03:40:59 /home/v-haipengluo/.conda/envs/llamax/bin/python -u src/train.py --local_rank=1 --model_name_or_path /workspaceblobstore/qins/test/20220316/haipeng/output_weights/llamax_13b_stackexchange_MATH_12w_sample_5w_score0.5_trainset_2e-5/checkpoint-992 --data_path /workspaceblobstore/qins/test/20220316/haipeng/data/Math_datasets/MATH_the_answer_is_format/hendrycks_math_7500_ori_gpt4_ori_15k.json --output_dir /workspaceblobstore/qins/test/20220316/haipeng/output_weights/llamax_13b_continue_train_stackMATH5w_checkpoint992_hendrycks_math_7500_ori_gpt4_ori_15k --num_train_epochs 3 --model_max_length 1150 --per_device_train_batch_size 17 --per_device_eval_batch_size 1 --gradient_accumulation_steps 1 --evaluation_strategy no --save_strategy steps --save_steps 36 --save_total_limit 200 --learning_rate 2e-5 --warmup_steps 10 --logging_steps 2 --lr_scheduler_type cosine --report_to tensorboard --gradient_checkpointing True --deepspeed src/configs/deepspeed_config.json --fp16 True
+v-haipe+ 18348 18333 99 05:34 pts/1    03:44:08 /home/v-haipengluo/.conda/envs/llamax/bin/python -u src/train.py --local_rank=2 --model_name_or_path /workspaceblobstore/qins/test/20220316/haipeng/output_weights/llamax_13b_stackexchange_MATH_12w_sample_5w_score0.5_trainset_2e-5/checkpoint-992 --data_path /workspaceblobstore/qins/test/20220316/haipeng/data/Math_datasets/MATH_the_answer_is_format/hendrycks_math_7500_ori_gpt4_ori_15k.json --output_dir /workspaceblobstore/qins/test/20220316/haipeng/output_weights/llamax_13b_continue_train_stackMATH5w_checkpoint992_hendrycks_math_7500_ori_gpt4_ori_15k --num_train_epochs 3 --model_max_length 1150 --per_device_train_batch_size 17 --per_device_eval_batch_size 1 --gradient_accumulation_steps 1 --evaluation_strategy no --save_strategy steps --save_steps 36 --save_total_limit 200 --learning_rate 2e-5 --warmup_steps 10 --logging_steps 2 --lr_scheduler_type cosine --report_to tensorboard --gradient_checkpointing True --deepspeed src/configs/deepspeed_config.json --fp16 True
+v-haipe+ 18349 18333 99 05:34 pts/1    03:32:51 /home/v-haipengluo/.conda/envs/llamax/bin/python -u src/train.py --local_rank=3 --model_name_or_path /workspaceblobstore/qins/test/20220316/haipeng/output_weights/llamax_13b_stackexchange_MATH_12w_sample_5w_score0.5_trainset_2e-5/checkpoint-992 --data_path /workspaceblobstore/qins/test/20220316/haipeng/data/Math_datasets/MATH_the_answer_is_format/hendrycks_math_7500_ori_gpt4_ori_15k.json --output_dir /workspaceblobstore/qins/test/20220316/haipeng/output_weights/llamax_13b_continue_train_stackMATH5w_checkpoint992_hendrycks_math_7500_ori_gpt4_ori_15k --num_train_epochs 3 --model_max_length 1150 --per_device_train_batch_size 17 --per_device_eval_batch_size 1 --gradient_accumulation_steps 1 --evaluation_strategy no --save_strategy steps --save_steps 36 --save_total_limit 200 --learning_rate 2e-5 --warmup_steps 10 --logging_steps 2 --lr_scheduler_type cosine --report_to tensorboard --gradient_checkpointing True --deepspeed src/configs/deepspeed_config.json --fp16 True
+v-haipe+ 18350 18333 99 05:34 pts/1    03:41:16 /home/v-haipengluo/.conda/envs/llamax/bin/python -u src/train.py --local_rank=4 --model_name_or_path /workspaceblobstore/qins/test/20220316/haipeng/output_weights/llamax_13b_stackexchange_MATH_12w_sample_5w_score0.5_trainset_2e-5/checkpoint-992 --data_path /workspaceblobstore/qins/test/20220316/haipeng/data/Math_datasets/MATH_the_answer_is_format/hendrycks_math_7500_ori_gpt4_ori_15k.json --output_dir /workspaceblobstore/qins/test/20220316/haipeng/output_weights/llamax_13b_continue_train_stackMATH5w_checkpoint992_hendrycks_math_7500_ori_gpt4_ori_15k --num_train_epochs 3 --model_max_length 1150 --per_device_train_batch_size 17 --per_device_eval_batch_size 1 --gradient_accumulation_steps 1 --evaluation_strategy no --save_strategy steps --save_steps 36 --save_total_limit 200 --learning_rate 2e-5 --warmup_steps 10 --logging_steps 2 --lr_scheduler_type cosine --report_to tensorboard --gradient_checkpointing True --deepspeed src/configs/deepspeed_config.json --fp16 True
+v-haipe+ 18351 18333 99 05:34 pts/1    03:42:27 /home/v-haipengluo/.conda/envs/llamax/bin/python -u src/train.py --local_rank=5 --model_name_or_path /workspaceblobstore/qins/test/20220316/haipeng/output_weights/llamax_13b_stackexchange_MATH_12w_sample_5w_score0.5_trainset_2e-5/checkpoint-992 --data_path /workspaceblobstore/qins/test/20220316/haipeng/data/Math_datasets/MATH_the_answer_is_format/hendrycks_math_7500_ori_gpt4_ori_15k.json --output_dir /workspaceblobstore/qins/test/20220316/haipeng/output_weights/llamax_13b_continue_train_stackMATH5w_checkpoint992_hendrycks_math_7500_ori_gpt4_ori_15k --num_train_epochs 3 --model_max_length 1150 --per_device_train_batch_size 17 --per_device_eval_batch_size 1 --gradient_accumulation_steps 1 --evaluation_strategy no --save_strategy steps --save_steps 36 --save_total_limit 200 --learning_rate 2e-5 --warmup_steps 10 --logging_steps 2 --lr_scheduler_type cosine --report_to tensorboard --gradient_checkpointing True --deepspeed src/configs/deepspeed_config.json --fp16 True
+v-haipe+ 24334 23818  0 Sep23 pts/7    00:00:25 python -m http.server

img/main.py

@@ -0,0 +1,528 @@
+import gc
+import math
+import multiprocessing
+import os
+import traceback
+from datetime import datetime
+from io import BytesIO
+from itertools import permutations
+from multiprocessing.pool import Pool
+from pathlib import Path
+from urllib.parse import quote_plus
+
+import numpy as np
+import nltk
+import torch
+
+from PIL.Image import Image
+from diffusers import DiffusionPipeline, StableDiffusionXLInpaintPipeline
+from diffusers.utils import load_image
+from fastapi import FastAPI
+from fastapi.middleware.gzip import GZipMiddleware
+from loguru import logger
+from starlette.middleware.cors import CORSMiddleware
+from starlette.responses import FileResponse
+from starlette.responses import JSONResponse
+
+from env import BUCKET_PATH, BUCKET_NAME
+# from stable_diffusion_server.bucket_api import check_if_blob_exists, upload_to_bucket
+torch._dynamo.config.suppress_errors = True
+
+import string
+import random
+
+def generate_save_path():
+    # initializing size of string
+    N = 7
+
+    # using random.choices()
+    # generating random strings
+    res = ''.join(random.choices(string.ascii_uppercase +
+                                 string.digits, k=N))
+    return res
+
+pipe = DiffusionPipeline.from_pretrained(
+    "models/stable-diffusion-xl-base-1.0",
+    torch_dtype=torch.bfloat16,
+    use_safetensors=True,
+    variant="fp16",
+    # safety_checker=None,
+)  # todo try torch_dtype=bfloat16
+pipe.watermark = None
+
+pipe.to("cuda")
+
+refiner = DiffusionPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-refiner-1.0",
+    text_encoder_2=pipe.text_encoder_2,
+    vae=pipe.vae,
+    torch_dtype=torch.bfloat16, # safer to use bfloat?
+    use_safetensors=True,
+    variant="fp16", #remember not to download the big model
+)
+refiner.watermark = None
+refiner.to("cuda")
+
+# {'scheduler', 'text_encoder', 'text_encoder_2', 'tokenizer', 'tokenizer_2', 'unet', 'vae'} can be passed in from existing model
+inpaintpipe = StableDiffusionXLInpaintPipeline.from_pretrained(
+    "models/stable-diffusion-xl-base-1.0", torch_dtype=torch.bfloat16, variant="fp16", use_safetensors=True,
+    scheduler=pipe.scheduler,
+    text_encoder=pipe.text_encoder,
+    text_encoder_2=pipe.text_encoder_2,
+    tokenizer=pipe.tokenizer,
+    tokenizer_2=pipe.tokenizer_2,
+    unet=pipe.unet,
+    vae=pipe.vae,
+    # load_connected_pipeline=
+)
+# # switch out to save gpu mem
+# del inpaintpipe.vae
+# del inpaintpipe.text_encoder_2
+# del inpaintpipe.text_encoder
+# del inpaintpipe.scheduler
+# del inpaintpipe.tokenizer
+# del inpaintpipe.tokenizer_2
+# del inpaintpipe.unet
+# inpaintpipe.vae = pipe.vae
+# inpaintpipe.text_encoder_2 = pipe.text_encoder_2
+# inpaintpipe.text_encoder = pipe.text_encoder
+# inpaintpipe.scheduler = pipe.scheduler
+# inpaintpipe.tokenizer = pipe.tokenizer
+# inpaintpipe.tokenizer_2 = pipe.tokenizer_2
+# inpaintpipe.unet = pipe.unet
+# todo this should work
+# inpaintpipe = StableDiffusionXLInpaintPipeline( # construct an inpainter using the existing model
+#     vae=pipe.vae,
+#     text_encoder_2=pipe.text_encoder_2,
+#     text_encoder=pipe.text_encoder,
+#     unet=pipe.unet,
+#     scheduler=pipe.scheduler,
+#     tokenizer=pipe.tokenizer,
+#     tokenizer_2=pipe.tokenizer_2,
+#     requires_aesthetics_score=False,
+# )
+inpaintpipe.to("cuda")
+inpaintpipe.watermark = None
+# inpaintpipe.register_to_config(requires_aesthetics_score=False)
+
+inpaint_refiner = StableDiffusionXLInpaintPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-refiner-1.0",
+    text_encoder_2=inpaintpipe.text_encoder_2,
+    vae=inpaintpipe.vae,
+    torch_dtype=torch.bfloat16,
+    use_safetensors=True,
+    variant="fp16",
+
+    tokenizer_2=refiner.tokenizer_2,
+    tokenizer=refiner.tokenizer,
+    scheduler=refiner.scheduler,
+    text_encoder=refiner.text_encoder,
+    unet=refiner.unet,
+)
+# del inpaint_refiner.vae
+# del inpaint_refiner.text_encoder_2
+# del inpaint_refiner.text_encoder
+# del inpaint_refiner.scheduler
+# del inpaint_refiner.tokenizer
+# del inpaint_refiner.tokenizer_2
+# del inpaint_refiner.unet
+# inpaint_refiner.vae = inpaintpipe.vae
+# inpaint_refiner.text_encoder_2 = inpaintpipe.text_encoder_2
+#
+# inpaint_refiner.text_encoder = refiner.text_encoder
+# inpaint_refiner.scheduler = refiner.scheduler
+# inpaint_refiner.tokenizer = refiner.tokenizer
+# inpaint_refiner.tokenizer_2 = refiner.tokenizer_2
+# inpaint_refiner.unet = refiner.unet
+
+# inpaint_refiner = StableDiffusionXLInpaintPipeline(
+#     text_encoder_2=inpaintpipe.text_encoder_2,
+#     vae=inpaintpipe.vae,
+#     # the rest from the existing refiner
+#     tokenizer_2=refiner.tokenizer_2,
+#     tokenizer=refiner.tokenizer,
+#     scheduler=refiner.scheduler,
+#     text_encoder=refiner.text_encoder,
+#     unet=refiner.unet,
+#     requires_aesthetics_score=False,
+# )
+inpaint_refiner.to("cuda")
+inpaint_refiner.watermark = None
+# inpaint_refiner.register_to_config(requires_aesthetics_score=False)
+
+n_steps = 40
+high_noise_frac = 0.8
+
+# if using torch < 2.0
+# pipe.enable_xformers_memory_efficient_attention()
+
+
+# pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
+# this can cause errors on some inputs so consider disabling it
+pipe.unet = torch.compile(pipe.unet)
+refiner.unet = torch.compile(refiner.unet)#, mode="reduce-overhead", fullgraph=True)
+# compile the inpainters - todo reuse the other unets? swap out the models for others/del them so they share models and can be swapped efficiently
+inpaintpipe.unet = pipe.unet
+inpaint_refiner.unet = refiner.unet
+# inpaintpipe.unet = torch.compile(inpaintpipe.unet)
+# inpaint_refiner.unet = torch.compile(inpaint_refiner.unet)
+from pydantic import BaseModel
+
+app = FastAPI(
+    openapi_url="/static/openapi.json",
+    docs_url="/swagger-docs",
+    redoc_url="/redoc",
+    title="Generate Images Netwrck API",
+    description="Character Chat API",
+    # root_path="https://api.text-generator.io",
+    version="1",
+)
+app.add_middleware(GZipMiddleware, minimum_size=1000)
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+stopwords = nltk.corpus.stopwords.words("english")
+
+class Img(BaseModel):
+    system_prompt: str
+    ASSISTANT: str
+
+# img_url = "http://phlrr2019.guest.corp.microsoft.com:8000/img1_sdv2.1.png"
+img_url = "http://phlrr3058.guest.corp.microsoft.com:8000/"#/img1_sdv2.1.png"
+
+@app.post("/image_url")
+def image_url(img: Img):
+    system_prompt = img.system_prompt
+    prompt = img.ASSISTANT
+    # if Path(save_path).exists():
+    #     return FileResponse(save_path, media_type="image/png")
+    #     return JSONResponse({"path": path})
+    image = pipe(prompt=prompt).images[0]
+    # if not save_path:
+    save_path = generate_save_path()
+    save_path = f"images/{save_path}.png"
+    image.save(save_path)
+    # save_path = '/'.join(path_components) + quote_plus(final_name)
+    path = f"{img_url}/{save_path}"
+    return JSONResponse({"path": path})
+
+
+@app.get("/make_image")
+# @app.post("/make_image")
+def make_image(prompt: str, save_path: str = ""):
+    if Path(save_path).exists():
+        return FileResponse(save_path, media_type="image/png")
+    image = pipe(prompt=prompt).images[0]
+    if not save_path:
+        save_path = f"images/{prompt}.png"
+    image.save(save_path)
+    return FileResponse(save_path, media_type="image/png")
+
+
+@app.get("/create_and_upload_image")
+def create_and_upload_image(prompt: str, width: int=1024, height:int=1024, save_path: str = ""):
+    path_components = save_path.split("/")[0:-1]
+    final_name = save_path.split("/")[-1]
+    if not path_components:
+        path_components = []
+    save_path = '/'.join(path_components) + quote_plus(final_name)
+    path = get_image_or_create_upload_to_cloud_storage(prompt, width, height, save_path)
+    return JSONResponse({"path": path})
+
+@app.get("/inpaint_and_upload_image")
+def inpaint_and_upload_image(prompt: str, image_url:str, mask_url:str, save_path: str = ""):
+    path_components = save_path.split("/")[0:-1]
+    final_name = save_path.split("/")[-1]
+    if not path_components:
+        path_components = []
+    save_path = '/'.join(path_components) + quote_plus(final_name)
+    path = get_image_or_inpaint_upload_to_cloud_storage(prompt, image_url, mask_url, save_path)
+    return JSONResponse({"path": path})
+
+
+def get_image_or_create_upload_to_cloud_storage(prompt:str,width:int, height:int, save_path:str):
+    prompt = shorten_too_long_text(prompt)
+    save_path = shorten_too_long_text(save_path)
+    # check exists - todo cache this
+    if check_if_blob_exists(save_path):
+        return f"https://{BUCKET_NAME}/{BUCKET_PATH}/{save_path}"
+    bio = create_image_from_prompt(prompt, width, height)
+    if bio is None:
+        return None # error thrown in pool
+    link = upload_to_bucket(save_path, bio, is_bytesio=True)
+    return link
+def get_image_or_inpaint_upload_to_cloud_storage(prompt:str, image_url:str, mask_url:str, save_path:str):
+    prompt = shorten_too_long_text(prompt)
+    save_path = shorten_too_long_text(save_path)
+    # check exists - todo cache this
+    if check_if_blob_exists(save_path):
+        return f"https://{BUCKET_NAME}/{BUCKET_PATH}/{save_path}"
+    bio = inpaint_image_from_prompt(prompt, image_url, mask_url)
+    if bio is None:
+        return None # error thrown in pool
+    link = upload_to_bucket(save_path, bio, is_bytesio=True)
+    return link
+
+# multiprocessing.set_start_method('spawn', True)
+# processes_pool = Pool(1) # cant do too much at once or OOM errors happen
+# def create_image_from_prompt_sync(prompt):
+#     """have to call this sync to avoid OOM errors"""
+#     return processes_pool.apply_async(create_image_from_prompt, args=(prompt,), ).wait()
+
+def create_image_from_prompt(prompt, width, height):
+    # round width and height down to multiple of 64
+    block_width = width - (width % 64)
+    block_height = height - (height % 64)
+    prompt = shorten_too_long_text(prompt)
+    # image = pipe(prompt=prompt).images[0]
+    try:
+        image = pipe(prompt=prompt,
+                     width=block_width,
+                     height=block_height,
+                     # denoising_end=high_noise_frac,
+                     # output_type='latent',
+                     # height=512,
+                     # width=512,
+                     num_inference_steps=50).images[0]  # normally uses 50 steps
+    except Exception as e:
+        # try rm stopwords + half the prompt
+        # todo try prompt permutations
+        logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+        prompt = ' '.join((word for word in prompt if word not in stopwords))
+        prompts = prompt.split()
+
+        prompt = ' '.join(prompts[:len(prompts) // 2])
+        logger.info(f"shortened prompt to: {len(prompt)}")
+        image = None
+        if prompt:
+            try:
+                image = pipe(prompt=prompt,
+                             width=block_width,
+                             height=block_height,
+                             # denoising_end=high_noise_frac,
+                             # output_type='latent',
+                             # height=512,
+                             # width=512,
+                             num_inference_steps=50).images[0]  # normally uses 50 steps
+            except Exception as e:
+                # logger.info("trying to permute prompt")
+                # # try two swaps of the prompt/permutations
+                # prompt = prompt.split()
+                # prompt = ' '.join(permutations(prompt, 2).__next__())
+                logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+                prompt = ' '.join((word for word in prompt if word not in stopwords))
+                prompts = prompt.split()
+
+                prompt = ' '.join(prompts[:len(prompts) // 2])
+                logger.info(f"shortened prompt to: {len(prompt)}")
+
+                try:
+                    image = pipe(prompt=prompt,
+                                 width=block_width,
+                                 height=block_height,
+                                 # denoising_end=high_noise_frac,
+                                 # output_type='latent', # dont need latent yet - we refine the image at full res
+                                 # height=512,
+                                 # width=512,
+                                 num_inference_steps=50).images[0]  # normally uses 50 steps
+                except Exception as e:
+                    # just error out
+                    traceback.print_exc()
+                    raise e
+                    # logger.info("restarting server to fix cuda issues (device side asserts)")
+                    # todo fix device side asserts instead of restart to fix
+                    # todo only restart the correct gunicorn
+                    # this could be really annoying if your running other gunicorns on your machine which also get restarted
+                    # os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+                    # os.system("kill -1 `pgrep gunicorn`")
+    # todo refine
+    # if image != None:
+    #     image = refiner(
+    #         prompt=prompt,
+    #         # width=block_width,
+    #         # height=block_height,
+    #         num_inference_steps=n_steps,
+    #         # denoising_start=high_noise_frac,
+    #         image=image,
+    #     ).images[0]
+    if width != block_width or height != block_height:
+        # resize to original size width/height
+        # find aspect ratio to scale up to that covers the original img input width/height
+        scale_up_ratio = max(width / block_width, height / block_height)
+        image = image.resize((math.ceil(block_width * scale_up_ratio), math.ceil(height * scale_up_ratio)))
+        # crop image to original size
+        image = image.crop((0, 0, width, height))
+    # try:
+    #     # gc.collect()
+    #     torch.cuda.empty_cache()
+    # except Exception as e:
+    #     traceback.print_exc()
+    #     logger.info("restarting server to fix cuda issues (device side asserts)")
+    #     # todo fix device side asserts instead of restart to fix
+    #     # todo only restart the correct gunicorn
+    #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+    #     os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+    #     os.system("kill -1 `pgrep gunicorn`")
+    # save as bytesio
+    bs = BytesIO()
+
+    bright_count = np.sum(np.array(image) > 0)
+    if bright_count == 0:
+        # we have a black image, this is an error likely we need a restart
+        logger.info("restarting server to fix cuda issues (device side asserts)")
+        #     # todo fix device side asserts instead of restart to fix
+        #     # todo only restart the correct gunicorn
+        #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+        os.system("kill -1 `pgrep gunicorn`")
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep uvicorn`")
+        os.system("kill -1 `pgrep uvicorn`")
+
+        return None
+    image.save(bs, quality=85, optimize=True, format="webp")
+    bio = bs.getvalue()
+    # touch progress.txt file - if we dont do this we get restarted by supervisor/other processes for reliability
+    with open("progress.txt", "w") as f:
+        current_time = datetime.now().strftime("%H:%M:%S")
+        f.write(f"{current_time}")
+    return bio
+
+def inpaint_image_from_prompt(prompt, image_url: str, mask_url: str):
+    prompt = shorten_too_long_text(prompt)
+    # image = pipe(prompt=prompt).images[0]
+
+    init_image = load_image(image_url).convert("RGB")
+    mask_image = load_image(mask_url).convert("RGB") # why rgb for a 1 channel mask?
+    num_inference_steps = 75
+    high_noise_frac = 0.7
+
+    try:
+        image = inpaintpipe(
+            prompt=prompt,
+            image=init_image,
+            mask_image=mask_image,
+            num_inference_steps=num_inference_steps,
+            denoising_start=high_noise_frac,
+            output_type="latent",
+        ).images[0]  # normally uses 50 steps
+    except Exception as e:
+        # try rm stopwords + half the prompt
+        # todo try prompt permutations
+        logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+        prompt = ' '.join((word for word in prompt if word not in stopwords))
+        prompts = prompt.split()
+
+        prompt = ' '.join(prompts[:len(prompts) // 2])
+        logger.info(f"shortened prompt to: {len(prompt)}")
+        image = None
+        if prompt:
+            try:
+                image = pipe(
+                    prompt=prompt,
+                    image=init_image,
+                    mask_image=mask_image,
+                    num_inference_steps=num_inference_steps,
+                    denoising_start=high_noise_frac,
+                    output_type="latent",
+                ).images[0]  # normally uses 50 steps
+            except Exception as e:
+                # logger.info("trying to permute prompt")
+                # # try two swaps of the prompt/permutations
+                # prompt = prompt.split()
+                # prompt = ' '.join(permutations(prompt, 2).__next__())
+                logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+                prompt = ' '.join((word for word in prompt if word not in stopwords))
+                prompts = prompt.split()
+
+                prompt = ' '.join(prompts[:len(prompts) // 2])
+                logger.info(f"shortened prompt to: {len(prompt)}")
+
+                try:
+                    image = inpaintpipe(
+                        prompt=prompt,
+                        image=init_image,
+                        mask_image=mask_image,
+                        num_inference_steps=num_inference_steps,
+                        denoising_start=high_noise_frac,
+                        output_type="latent",
+                    ).images[0]  # normally uses 50 steps
+                except Exception as e:
+                    # just error out
+                    traceback.print_exc()
+                    raise e
+                    # logger.info("restarting server to fix cuda issues (device side asserts)")
+                    # todo fix device side asserts instead of restart to fix
+                    # todo only restart the correct gunicorn
+                    # this could be really annoying if your running other gunicorns on your machine which also get restarted
+                    # os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+                    # os.system("kill -1 `pgrep gunicorn`")
+    if image != None:
+        image = inpaint_refiner(
+            prompt=prompt,
+            image=image,
+            mask_image=mask_image,
+            num_inference_steps=num_inference_steps,
+            denoising_start=high_noise_frac,
+
+        ).images[0]
+    # try:
+    #     # gc.collect()
+    #     torch.cuda.empty_cache()
+    # except Exception as e:
+    #     traceback.print_exc()
+    #     logger.info("restarting server to fix cuda issues (device side asserts)")
+    #     # todo fix device side asserts instead of restart to fix
+    #     # todo only restart the correct gunicorn
+    #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+    #     os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+    #     os.system("kill -1 `pgrep gunicorn`")
+    # save as bytesio
+    bs = BytesIO()
+
+    bright_count = np.sum(np.array(image) > 0)
+    if bright_count == 0:
+        # we have a black image, this is an error likely we need a restart
+        logger.info("restarting server to fix cuda issues (device side asserts)")
+        #     # todo fix device side asserts instead of restart to fix
+        #     # todo only restart the correct gunicorn
+        #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+        os.system("kill -1 `pgrep gunicorn`")
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep uvicorn`")
+        os.system("kill -1 `pgrep uvicorn`")
+
+        return None
+    image.save(bs, quality=85, optimize=True, format="webp")
+    bio = bs.getvalue()
+    # touch progress.txt file - if we dont do this we get restarted by supervisor/other processes for reliability
+    with open("progress.txt", "w") as f:
+        current_time = datetime.now().strftime("%H:%M:%S")
+        f.write(f"{current_time}")
+    return bio
+
+
+
+def shorten_too_long_text(prompt):
+    if len(prompt) > 200:
+        # remove stopwords
+        prompt = prompt.split() # todo also split hyphens
+        prompt = ' '.join((word for word in prompt if word not in stopwords))
+        if len(prompt) > 200:
+            prompt = prompt[:200]
+    return prompt
+
+# image = pipe(prompt=prompt).images[0]
+#
+# image.save("test.png")
+# # save all images
+# for i, image in enumerate(images):
+#     image.save(f"{i}.png")

img/main_1024.py

@@ -0,0 +1,549 @@
+import gc
+import math
+import multiprocessing
+import os
+import traceback
+from datetime import datetime
+from io import BytesIO
+from itertools import permutations
+from multiprocessing.pool import Pool
+from pathlib import Path
+from urllib.parse import quote_plus
+
+import numpy as np
+import nltk
+import torch
+
+from PIL.Image import Image
+from diffusers import DiffusionPipeline, StableDiffusionXLInpaintPipeline
+from diffusers.utils import load_image
+from fastapi import FastAPI
+from fastapi.middleware.gzip import GZipMiddleware
+from loguru import logger
+from starlette.middleware.cors import CORSMiddleware
+from starlette.responses import FileResponse
+from starlette.responses import JSONResponse
+
+from env import BUCKET_PATH, BUCKET_NAME
+# from stable_diffusion_server.bucket_api import check_if_blob_exists, upload_to_bucket
+torch._dynamo.config.suppress_errors = True
+
+import string
+import random
+
+def generate_save_path():
+    # initializing size of string
+    N = 7
+
+    # using random.choices()
+    # generating random strings
+    res = ''.join(random.choices(string.ascii_uppercase +
+                                 string.digits, k=N))
+    return res
+
+# pipe = DiffusionPipeline.from_pretrained(
+#     "models/stable-diffusion-xl-base-1.0",
+#     torch_dtype=torch.bfloat16,
+#     use_safetensors=True,
+#     variant="fp16",
+#     # safety_checker=None,
+# )  # todo try torch_dtype=bfloat16
+
+model_dir = os.getenv("SDXL_MODEL_DIR")
+
+if model_dir:
+    # Use local model
+    model_key_base = os.path.join(model_dir, "stable-diffusion-xl-base-1.0")
+    model_key_refiner = os.path.join(model_dir, "stable-diffusion-xl-refiner-1.0")
+else:
+    model_key_base = "stabilityai/stable-diffusion-xl-base-1.0"
+    model_key_refiner = "stabilityai/stable-diffusion-xl-refiner-1.0"
+
+pipe = DiffusionPipeline.from_pretrained(model_key_base, torch_dtype=torch.float16, use_safetensors=True, variant="fp16")
+
+pipe.watermark = None
+
+pipe.to("cuda")
+
+refiner = DiffusionPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-refiner-1.0",
+    text_encoder_2=pipe.text_encoder_2,
+    vae=pipe.vae,
+    torch_dtype=torch.bfloat16, # safer to use bfloat?
+    use_safetensors=True,
+    variant="fp16", #remember not to download the big model
+)
+refiner.watermark = None
+refiner.to("cuda")
+
+# {'scheduler', 'text_encoder', 'text_encoder_2', 'tokenizer', 'tokenizer_2', 'unet', 'vae'} can be passed in from existing model
+inpaintpipe = StableDiffusionXLInpaintPipeline.from_pretrained(
+    "models/stable-diffusion-xl-base-1.0", torch_dtype=torch.bfloat16, variant="fp16", use_safetensors=True,
+    scheduler=pipe.scheduler,
+    text_encoder=pipe.text_encoder,
+    text_encoder_2=pipe.text_encoder_2,
+    tokenizer=pipe.tokenizer,
+    tokenizer_2=pipe.tokenizer_2,
+    unet=pipe.unet,
+    vae=pipe.vae,
+    # load_connected_pipeline=
+)
+# # switch out to save gpu mem
+# del inpaintpipe.vae
+# del inpaintpipe.text_encoder_2
+# del inpaintpipe.text_encoder
+# del inpaintpipe.scheduler
+# del inpaintpipe.tokenizer
+# del inpaintpipe.tokenizer_2
+# del inpaintpipe.unet
+# inpaintpipe.vae = pipe.vae
+# inpaintpipe.text_encoder_2 = pipe.text_encoder_2
+# inpaintpipe.text_encoder = pipe.text_encoder
+# inpaintpipe.scheduler = pipe.scheduler
+# inpaintpipe.tokenizer = pipe.tokenizer
+# inpaintpipe.tokenizer_2 = pipe.tokenizer_2
+# inpaintpipe.unet = pipe.unet
+# todo this should work
+# inpaintpipe = StableDiffusionXLInpaintPipeline( # construct an inpainter using the existing model
+#     vae=pipe.vae,
+#     text_encoder_2=pipe.text_encoder_2,
+#     text_encoder=pipe.text_encoder,
+#     unet=pipe.unet,
+#     scheduler=pipe.scheduler,
+#     tokenizer=pipe.tokenizer,
+#     tokenizer_2=pipe.tokenizer_2,
+#     requires_aesthetics_score=False,
+# )
+inpaintpipe.to("cuda")
+inpaintpipe.watermark = None
+# inpaintpipe.register_to_config(requires_aesthetics_score=False)
+
+inpaint_refiner = StableDiffusionXLInpaintPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-refiner-1.0",
+    text_encoder_2=inpaintpipe.text_encoder_2,
+    vae=inpaintpipe.vae,
+    torch_dtype=torch.bfloat16,
+    use_safetensors=True,
+    variant="fp16",
+
+    tokenizer_2=refiner.tokenizer_2,
+    tokenizer=refiner.tokenizer,
+    scheduler=refiner.scheduler,
+    text_encoder=refiner.text_encoder,
+    unet=refiner.unet,
+)
+# del inpaint_refiner.vae
+# del inpaint_refiner.text_encoder_2
+# del inpaint_refiner.text_encoder
+# del inpaint_refiner.scheduler
+# del inpaint_refiner.tokenizer
+# del inpaint_refiner.tokenizer_2
+# del inpaint_refiner.unet
+# inpaint_refiner.vae = inpaintpipe.vae
+# inpaint_refiner.text_encoder_2 = inpaintpipe.text_encoder_2
+#
+# inpaint_refiner.text_encoder = refiner.text_encoder
+# inpaint_refiner.scheduler = refiner.scheduler
+# inpaint_refiner.tokenizer = refiner.tokenizer
+# inpaint_refiner.tokenizer_2 = refiner.tokenizer_2
+# inpaint_refiner.unet = refiner.unet
+
+# inpaint_refiner = StableDiffusionXLInpaintPipeline(
+#     text_encoder_2=inpaintpipe.text_encoder_2,
+#     vae=inpaintpipe.vae,
+#     # the rest from the existing refiner
+#     tokenizer_2=refiner.tokenizer_2,
+#     tokenizer=refiner.tokenizer,
+#     scheduler=refiner.scheduler,
+#     text_encoder=refiner.text_encoder,
+#     unet=refiner.unet,
+#     requires_aesthetics_score=False,
+# )
+inpaint_refiner.to("cuda")
+inpaint_refiner.watermark = None
+# inpaint_refiner.register_to_config(requires_aesthetics_score=False)
+
+n_steps = 40
+high_noise_frac = 0.8
+
+# if using torch < 2.0
+# pipe.enable_xformers_memory_efficient_attention()
+
+
+# pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
+# this can cause errors on some inputs so consider disabling it
+pipe.unet = torch.compile(pipe.unet)
+refiner.unet = torch.compile(refiner.unet)#, mode="reduce-overhead", fullgraph=True)
+# compile the inpainters - todo reuse the other unets? swap out the models for others/del them so they share models and can be swapped efficiently
+inpaintpipe.unet = pipe.unet
+inpaint_refiner.unet = refiner.unet
+# inpaintpipe.unet = torch.compile(inpaintpipe.unet)
+# inpaint_refiner.unet = torch.compile(inpaint_refiner.unet)
+from pydantic import BaseModel
+
+app = FastAPI(
+    openapi_url="/static/openapi.json",
+    docs_url="/swagger-docs",
+    redoc_url="/redoc",
+    title="Generate Images Netwrck API",
+    description="Character Chat API",
+    # root_path="https://api.text-generator.io",
+    version="1",
+)
+app.add_middleware(GZipMiddleware, minimum_size=1000)
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+stopwords = nltk.corpus.stopwords.words("english")
+
+class Img(BaseModel):
+    system_prompt: str
+    ASSISTANT: str
+
+# img_url = "http://phlrr2019.guest.corp.microsoft.com:8000/img1_sdv2.1.png"
+img_url = "http://phlrr3058.guest.corp.microsoft.com:8000/"#/img1_sdv2.1.png"
+
+is_gpu_busy = False
+
+
+@app.post("/image_url")
+def image_url(img: Img):
+    system_prompt = img.system_prompt
+    prompt = img.ASSISTANT
+    # if Path(save_path).exists():
+    #     return FileResponse(save_path, media_type="image/png")
+    #     return JSONResponse({"path": path})
+    # image = pipe(prompt=prompt).images[0]
+    g = torch.Generator(device="cuda")
+    # image = pipe(prompt=prompt, width=1024, height=1024, generator=g).images[0]
+    image = pipe(prompt=prompt, width=1024, height=1024).images[0]
+
+    # if not save_path:
+    save_path = generate_save_path()
+    save_path = f"images/{save_path}.png"
+    image.save(save_path)
+    # save_path = '/'.join(path_components) + quote_plus(final_name)
+    path = f"{img_url}/{save_path}"
+    return JSONResponse({"path": path})
+
+
+@app.get("/make_image")
+# @app.post("/make_image")
+def make_image(prompt: str, save_path: str = ""):
+    if Path(save_path).exists():
+        return FileResponse(save_path, media_type="image/png")
+    image = pipe(prompt=prompt).images[0]
+    if not save_path:
+        save_path = f"images/{prompt}.png"
+    image.save(save_path)
+    return FileResponse(save_path, media_type="image/png")
+
+
+@app.get("/create_and_upload_image")
+def create_and_upload_image(prompt: str, width: int=1024, height:int=1024, save_path: str = ""):
+    path_components = save_path.split("/")[0:-1]
+    final_name = save_path.split("/")[-1]
+    if not path_components:
+        path_components = []
+    save_path = '/'.join(path_components) + quote_plus(final_name)
+    path = get_image_or_create_upload_to_cloud_storage(prompt, width, height, save_path)
+    return JSONResponse({"path": path})
+
+@app.get("/inpaint_and_upload_image")
+def inpaint_and_upload_image(prompt: str, image_url:str, mask_url:str, save_path: str = ""):
+    path_components = save_path.split("/")[0:-1]
+    final_name = save_path.split("/")[-1]
+    if not path_components:
+        path_components = []
+    save_path = '/'.join(path_components) + quote_plus(final_name)
+    path = get_image_or_inpaint_upload_to_cloud_storage(prompt, image_url, mask_url, save_path)
+    return JSONResponse({"path": path})
+
+
+def get_image_or_create_upload_to_cloud_storage(prompt:str,width:int, height:int, save_path:str):
+    prompt = shorten_too_long_text(prompt)
+    save_path = shorten_too_long_text(save_path)
+    # check exists - todo cache this
+    if check_if_blob_exists(save_path):
+        return f"https://{BUCKET_NAME}/{BUCKET_PATH}/{save_path}"
+    bio = create_image_from_prompt(prompt, width, height)
+    if bio is None:
+        return None # error thrown in pool
+    link = upload_to_bucket(save_path, bio, is_bytesio=True)
+    return link
+def get_image_or_inpaint_upload_to_cloud_storage(prompt:str, image_url:str, mask_url:str, save_path:str):
+    prompt = shorten_too_long_text(prompt)
+    save_path = shorten_too_long_text(save_path)
+    # check exists - todo cache this
+    if check_if_blob_exists(save_path):
+        return f"https://{BUCKET_NAME}/{BUCKET_PATH}/{save_path}"
+    bio = inpaint_image_from_prompt(prompt, image_url, mask_url)
+    if bio is None:
+        return None # error thrown in pool
+    link = upload_to_bucket(save_path, bio, is_bytesio=True)
+    return link
+
+# multiprocessing.set_start_method('spawn', True)
+# processes_pool = Pool(1) # cant do too much at once or OOM errors happen
+# def create_image_from_prompt_sync(prompt):
+#     """have to call this sync to avoid OOM errors"""
+#     return processes_pool.apply_async(create_image_from_prompt, args=(prompt,), ).wait()
+
+def create_image_from_prompt(prompt, width, height):
+    # round width and height down to multiple of 64
+    block_width = width - (width % 64)
+    block_height = height - (height % 64)
+    prompt = shorten_too_long_text(prompt)
+    # image = pipe(prompt=prompt).images[0]
+    try:
+        image = pipe(prompt=prompt,
+                     width=block_width,
+                     height=block_height,
+                     # denoising_end=high_noise_frac,
+                     # output_type='latent',
+                     # height=512,
+                     # width=512,
+                     num_inference_steps=50).images[0]  # normally uses 50 steps
+    except Exception as e:
+        # try rm stopwords + half the prompt
+        # todo try prompt permutations
+        logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+        prompt = ' '.join((word for word in prompt if word not in stopwords))
+        prompts = prompt.split()
+
+        prompt = ' '.join(prompts[:len(prompts) // 2])
+        logger.info(f"shortened prompt to: {len(prompt)}")
+        image = None
+        if prompt:
+            try:
+                image = pipe(prompt=prompt,
+                             width=block_width,
+                             height=block_height,
+                             # denoising_end=high_noise_frac,
+                             # output_type='latent',
+                             # height=512,
+                             # width=512,
+                             num_inference_steps=50).images[0]  # normally uses 50 steps
+            except Exception as e:
+                # logger.info("trying to permute prompt")
+                # # try two swaps of the prompt/permutations
+                # prompt = prompt.split()
+                # prompt = ' '.join(permutations(prompt, 2).__next__())
+                logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+                prompt = ' '.join((word for word in prompt if word not in stopwords))
+                prompts = prompt.split()
+
+                prompt = ' '.join(prompts[:len(prompts) // 2])
+                logger.info(f"shortened prompt to: {len(prompt)}")
+
+                try:
+                    image = pipe(prompt=prompt,
+                                 width=block_width,
+                                 height=block_height,
+                                 # denoising_end=high_noise_frac,
+                                 # output_type='latent', # dont need latent yet - we refine the image at full res
+                                 # height=512,
+                                 # width=512,
+                                 num_inference_steps=50).images[0]  # normally uses 50 steps
+                except Exception as e:
+                    # just error out
+                    traceback.print_exc()
+                    raise e
+                    # logger.info("restarting server to fix cuda issues (device side asserts)")
+                    # todo fix device side asserts instead of restart to fix
+                    # todo only restart the correct gunicorn
+                    # this could be really annoying if your running other gunicorns on your machine which also get restarted
+                    # os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+                    # os.system("kill -1 `pgrep gunicorn`")
+    # todo refine
+    # if image != None:
+    #     image = refiner(
+    #         prompt=prompt,
+    #         # width=block_width,
+    #         # height=block_height,
+    #         num_inference_steps=n_steps,
+    #         # denoising_start=high_noise_frac,
+    #         image=image,
+    #     ).images[0]
+    if width != block_width or height != block_height:
+        # resize to original size width/height
+        # find aspect ratio to scale up to that covers the original img input width/height
+        scale_up_ratio = max(width / block_width, height / block_height)
+        image = image.resize((math.ceil(block_width * scale_up_ratio), math.ceil(height * scale_up_ratio)))
+        # crop image to original size
+        image = image.crop((0, 0, width, height))
+    # try:
+    #     # gc.collect()
+    #     torch.cuda.empty_cache()
+    # except Exception as e:
+    #     traceback.print_exc()
+    #     logger.info("restarting server to fix cuda issues (device side asserts)")
+    #     # todo fix device side asserts instead of restart to fix
+    #     # todo only restart the correct gunicorn
+    #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+    #     os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+    #     os.system("kill -1 `pgrep gunicorn`")
+    # save as bytesio
+    bs = BytesIO()
+
+    bright_count = np.sum(np.array(image) > 0)
+    if bright_count == 0:
+        # we have a black image, this is an error likely we need a restart
+        logger.info("restarting server to fix cuda issues (device side asserts)")
+        #     # todo fix device side asserts instead of restart to fix
+        #     # todo only restart the correct gunicorn
+        #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+        os.system("kill -1 `pgrep gunicorn`")
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep uvicorn`")
+        os.system("kill -1 `pgrep uvicorn`")
+
+        return None
+    image.save(bs, quality=85, optimize=True, format="webp")
+    bio = bs.getvalue()
+    # touch progress.txt file - if we dont do this we get restarted by supervisor/other processes for reliability
+    with open("progress.txt", "w") as f:
+        current_time = datetime.now().strftime("%H:%M:%S")
+        f.write(f"{current_time}")
+    return bio
+
+def inpaint_image_from_prompt(prompt, image_url: str, mask_url: str):
+    prompt = shorten_too_long_text(prompt)
+    # image = pipe(prompt=prompt).images[0]
+
+    init_image = load_image(image_url).convert("RGB")
+    mask_image = load_image(mask_url).convert("RGB") # why rgb for a 1 channel mask?
+    num_inference_steps = 75
+    high_noise_frac = 0.7
+
+    try:
+        image = inpaintpipe(
+            prompt=prompt,
+            image=init_image,
+            mask_image=mask_image,
+            num_inference_steps=num_inference_steps,
+            denoising_start=high_noise_frac,
+            output_type="latent",
+        ).images[0]  # normally uses 50 steps
+    except Exception as e:
+        # try rm stopwords + half the prompt
+        # todo try prompt permutations
+        logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+        prompt = ' '.join((word for word in prompt if word not in stopwords))
+        prompts = prompt.split()
+
+        prompt = ' '.join(prompts[:len(prompts) // 2])
+        logger.info(f"shortened prompt to: {len(prompt)}")
+        image = None
+        if prompt:
+            try:
+                image = pipe(
+                    prompt=prompt,
+                    image=init_image,
+                    mask_image=mask_image,
+                    num_inference_steps=num_inference_steps,
+                    denoising_start=high_noise_frac,
+                    output_type="latent",
+                ).images[0]  # normally uses 50 steps
+            except Exception as e:
+                # logger.info("trying to permute prompt")
+                # # try two swaps of the prompt/permutations
+                # prompt = prompt.split()
+                # prompt = ' '.join(permutations(prompt, 2).__next__())
+                logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+                prompt = ' '.join((word for word in prompt if word not in stopwords))
+                prompts = prompt.split()
+
+                prompt = ' '.join(prompts[:len(prompts) // 2])
+                logger.info(f"shortened prompt to: {len(prompt)}")
+
+                try:
+                    image = inpaintpipe(
+                        prompt=prompt,
+                        image=init_image,
+                        mask_image=mask_image,
+                        num_inference_steps=num_inference_steps,
+                        denoising_start=high_noise_frac,
+                        output_type="latent",
+                    ).images[0]  # normally uses 50 steps
+                except Exception as e:
+                    # just error out
+                    traceback.print_exc()
+                    raise e
+                    # logger.info("restarting server to fix cuda issues (device side asserts)")
+                    # todo fix device side asserts instead of restart to fix
+                    # todo only restart the correct gunicorn
+                    # this could be really annoying if your running other gunicorns on your machine which also get restarted
+                    # os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+                    # os.system("kill -1 `pgrep gunicorn`")
+    if image != None:
+        image = inpaint_refiner(
+            prompt=prompt,
+            image=image,
+            mask_image=mask_image,
+            num_inference_steps=num_inference_steps,
+            denoising_start=high_noise_frac,
+
+        ).images[0]
+    # try:
+    #     # gc.collect()
+    #     torch.cuda.empty_cache()
+    # except Exception as e:
+    #     traceback.print_exc()
+    #     logger.info("restarting server to fix cuda issues (device side asserts)")
+    #     # todo fix device side asserts instead of restart to fix
+    #     # todo only restart the correct gunicorn
+    #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+    #     os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+    #     os.system("kill -1 `pgrep gunicorn`")
+    # save as bytesio
+    bs = BytesIO()
+
+    bright_count = np.sum(np.array(image) > 0)
+    if bright_count == 0:
+        # we have a black image, this is an error likely we need a restart
+        logger.info("restarting server to fix cuda issues (device side asserts)")
+        #     # todo fix device side asserts instead of restart to fix
+        #     # todo only restart the correct gunicorn
+        #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+        os.system("kill -1 `pgrep gunicorn`")
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep uvicorn`")
+        os.system("kill -1 `pgrep uvicorn`")
+
+        return None
+    image.save(bs, quality=85, optimize=True, format="webp")
+    bio = bs.getvalue()
+    # touch progress.txt file - if we dont do this we get restarted by supervisor/other processes for reliability
+    with open("progress.txt", "w") as f:
+        current_time = datetime.now().strftime("%H:%M:%S")
+        f.write(f"{current_time}")
+    return bio
+
+
+
+def shorten_too_long_text(prompt):
+    if len(prompt) > 200:
+        # remove stopwords
+        prompt = prompt.split() # todo also split hyphens
+        prompt = ' '.join((word for word in prompt if word not in stopwords))
+        if len(prompt) > 200:
+            prompt = prompt[:200]
+    return prompt
+
+# image = pipe(prompt=prompt).images[0]
+#
+# image.save("test.png")
+# # save all images
+# for i, image in enumerate(images):
+#     image.save(f"{i}.png")
+

img/main_v2.py

@@ -0,0 +1,548 @@
+import gc
+import math
+import multiprocessing
+import os
+import traceback
+from datetime import datetime
+from io import BytesIO
+from itertools import permutations
+from multiprocessing.pool import Pool
+from pathlib import Path
+from urllib.parse import quote_plus
+
+import numpy as np
+import nltk
+import torch
+
+from PIL.Image import Image
+from diffusers import DiffusionPipeline, StableDiffusionXLInpaintPipeline
+from diffusers.utils import load_image
+from fastapi import FastAPI
+from fastapi.middleware.gzip import GZipMiddleware
+from loguru import logger
+from starlette.middleware.cors import CORSMiddleware
+from starlette.responses import FileResponse
+from starlette.responses import JSONResponse
+
+from env import BUCKET_PATH, BUCKET_NAME
+# from stable_diffusion_server.bucket_api import check_if_blob_exists, upload_to_bucket
+torch._dynamo.config.suppress_errors = True
+
+import string
+import random
+
+def generate_save_path():
+    # initializing size of string
+    N = 7
+
+    # using random.choices()
+    # generating random strings
+    res = ''.join(random.choices(string.ascii_uppercase +
+                                 string.digits, k=N))
+    return res
+
+# pipe = DiffusionPipeline.from_pretrained(
+#     "models/stable-diffusion-xl-base-1.0",
+#     torch_dtype=torch.bfloat16,
+#     use_safetensors=True,
+#     variant="fp16",
+#     # safety_checker=None,
+# )  # todo try torch_dtype=bfloat16
+
+model_dir = os.getenv("SDXL_MODEL_DIR")
+
+if model_dir:
+    # Use local model
+    model_key_base = os.path.join(model_dir, "stable-diffusion-xl-base-1.0")
+    model_key_refiner = os.path.join(model_dir, "stable-diffusion-xl-refiner-1.0")
+else:
+    model_key_base = "stabilityai/stable-diffusion-xl-base-1.0"
+    model_key_refiner = "stabilityai/stable-diffusion-xl-refiner-1.0"
+
+pipe = DiffusionPipeline.from_pretrained(model_key_base, torch_dtype=torch.float16, use_safetensors=True, variant="fp16")
+
+pipe.watermark = None
+
+pipe.to("cuda")
+
+refiner = DiffusionPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-refiner-1.0",
+    text_encoder_2=pipe.text_encoder_2,
+    vae=pipe.vae,
+    torch_dtype=torch.bfloat16, # safer to use bfloat?
+    use_safetensors=True,
+    variant="fp16", #remember not to download the big model
+)
+refiner.watermark = None
+refiner.to("cuda")
+
+# {'scheduler', 'text_encoder', 'text_encoder_2', 'tokenizer', 'tokenizer_2', 'unet', 'vae'} can be passed in from existing model
+inpaintpipe = StableDiffusionXLInpaintPipeline.from_pretrained(
+    "models/stable-diffusion-xl-base-1.0", torch_dtype=torch.bfloat16, variant="fp16", use_safetensors=True,
+    scheduler=pipe.scheduler,
+    text_encoder=pipe.text_encoder,
+    text_encoder_2=pipe.text_encoder_2,
+    tokenizer=pipe.tokenizer,
+    tokenizer_2=pipe.tokenizer_2,
+    unet=pipe.unet,
+    vae=pipe.vae,
+    # load_connected_pipeline=
+)
+# # switch out to save gpu mem
+# del inpaintpipe.vae
+# del inpaintpipe.text_encoder_2
+# del inpaintpipe.text_encoder
+# del inpaintpipe.scheduler
+# del inpaintpipe.tokenizer
+# del inpaintpipe.tokenizer_2
+# del inpaintpipe.unet
+# inpaintpipe.vae = pipe.vae
+# inpaintpipe.text_encoder_2 = pipe.text_encoder_2
+# inpaintpipe.text_encoder = pipe.text_encoder
+# inpaintpipe.scheduler = pipe.scheduler
+# inpaintpipe.tokenizer = pipe.tokenizer
+# inpaintpipe.tokenizer_2 = pipe.tokenizer_2
+# inpaintpipe.unet = pipe.unet
+# todo this should work
+# inpaintpipe = StableDiffusionXLInpaintPipeline( # construct an inpainter using the existing model
+#     vae=pipe.vae,
+#     text_encoder_2=pipe.text_encoder_2,
+#     text_encoder=pipe.text_encoder,
+#     unet=pipe.unet,
+#     scheduler=pipe.scheduler,
+#     tokenizer=pipe.tokenizer,
+#     tokenizer_2=pipe.tokenizer_2,
+#     requires_aesthetics_score=False,
+# )
+inpaintpipe.to("cuda")
+inpaintpipe.watermark = None
+# inpaintpipe.register_to_config(requires_aesthetics_score=False)
+
+inpaint_refiner = StableDiffusionXLInpaintPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-refiner-1.0",
+    text_encoder_2=inpaintpipe.text_encoder_2,
+    vae=inpaintpipe.vae,
+    torch_dtype=torch.bfloat16,
+    use_safetensors=True,
+    variant="fp16",
+
+    tokenizer_2=refiner.tokenizer_2,
+    tokenizer=refiner.tokenizer,
+    scheduler=refiner.scheduler,
+    text_encoder=refiner.text_encoder,
+    unet=refiner.unet,
+)
+# del inpaint_refiner.vae
+# del inpaint_refiner.text_encoder_2
+# del inpaint_refiner.text_encoder
+# del inpaint_refiner.scheduler
+# del inpaint_refiner.tokenizer
+# del inpaint_refiner.tokenizer_2
+# del inpaint_refiner.unet
+# inpaint_refiner.vae = inpaintpipe.vae
+# inpaint_refiner.text_encoder_2 = inpaintpipe.text_encoder_2
+#
+# inpaint_refiner.text_encoder = refiner.text_encoder
+# inpaint_refiner.scheduler = refiner.scheduler
+# inpaint_refiner.tokenizer = refiner.tokenizer
+# inpaint_refiner.tokenizer_2 = refiner.tokenizer_2
+# inpaint_refiner.unet = refiner.unet
+
+# inpaint_refiner = StableDiffusionXLInpaintPipeline(
+#     text_encoder_2=inpaintpipe.text_encoder_2,
+#     vae=inpaintpipe.vae,
+#     # the rest from the existing refiner
+#     tokenizer_2=refiner.tokenizer_2,
+#     tokenizer=refiner.tokenizer,
+#     scheduler=refiner.scheduler,
+#     text_encoder=refiner.text_encoder,
+#     unet=refiner.unet,
+#     requires_aesthetics_score=False,
+# )
+inpaint_refiner.to("cuda")
+inpaint_refiner.watermark = None
+# inpaint_refiner.register_to_config(requires_aesthetics_score=False)
+
+n_steps = 40
+high_noise_frac = 0.8
+
+# if using torch < 2.0
+# pipe.enable_xformers_memory_efficient_attention()
+
+
+# pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
+# this can cause errors on some inputs so consider disabling it
+pipe.unet = torch.compile(pipe.unet)
+refiner.unet = torch.compile(refiner.unet)#, mode="reduce-overhead", fullgraph=True)
+# compile the inpainters - todo reuse the other unets? swap out the models for others/del them so they share models and can be swapped efficiently
+inpaintpipe.unet = pipe.unet
+inpaint_refiner.unet = refiner.unet
+# inpaintpipe.unet = torch.compile(inpaintpipe.unet)
+# inpaint_refiner.unet = torch.compile(inpaint_refiner.unet)
+from pydantic import BaseModel
+
+app = FastAPI(
+    openapi_url="/static/openapi.json",
+    docs_url="/swagger-docs",
+    redoc_url="/redoc",
+    title="Generate Images Netwrck API",
+    description="Character Chat API",
+    # root_path="https://api.text-generator.io",
+    version="1",
+)
+app.add_middleware(GZipMiddleware, minimum_size=1000)
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+stopwords = nltk.corpus.stopwords.words("english")
+
+class Img(BaseModel):
+    system_prompt: str
+    ASSISTANT: str
+
+# img_url = "http://phlrr2019.guest.corp.microsoft.com:8000/img1_sdv2.1.png"
+img_url = "http://phlrr3105.guest.corp.microsoft.com:8000/"#/img1_sdv2.1.png"
+
+is_gpu_busy = False
+
+
+@app.post("/image_url")
+def image_url(img: Img):
+    system_prompt = img.system_prompt
+    prompt = img.ASSISTANT
+    # if Path(save_path).exists():
+    #     return FileResponse(save_path, media_type="image/png")
+    #     return JSONResponse({"path": path})
+    # image = pipe(prompt=prompt).images[0]
+    g = torch.Generator(device="cuda")
+    image = pipe(prompt=prompt, width=1024, height=1024, generator=g).images[0]
+
+    # if not save_path:
+    save_path = generate_save_path()
+    save_path = f"images/{save_path}.png"
+    image.save(save_path)
+    # save_path = '/'.join(path_components) + quote_plus(final_name)
+    path = f"{img_url}/{save_path}"
+    return JSONResponse({"path": path})
+
+
+@app.get("/make_image")
+# @app.post("/make_image")
+def make_image(prompt: str, save_path: str = ""):
+    if Path(save_path).exists():
+        return FileResponse(save_path, media_type="image/png")
+    image = pipe(prompt=prompt).images[0]
+    if not save_path:
+        save_path = f"images/{prompt}.png"
+    image.save(save_path)
+    return FileResponse(save_path, media_type="image/png")
+
+
+@app.get("/create_and_upload_image")
+def create_and_upload_image(prompt: str, width: int=1024, height:int=1024, save_path: str = ""):
+    path_components = save_path.split("/")[0:-1]
+    final_name = save_path.split("/")[-1]
+    if not path_components:
+        path_components = []
+    save_path = '/'.join(path_components) + quote_plus(final_name)
+    path = get_image_or_create_upload_to_cloud_storage(prompt, width, height, save_path)
+    return JSONResponse({"path": path})
+
+@app.get("/inpaint_and_upload_image")
+def inpaint_and_upload_image(prompt: str, image_url:str, mask_url:str, save_path: str = ""):
+    path_components = save_path.split("/")[0:-1]
+    final_name = save_path.split("/")[-1]
+    if not path_components:
+        path_components = []
+    save_path = '/'.join(path_components) + quote_plus(final_name)
+    path = get_image_or_inpaint_upload_to_cloud_storage(prompt, image_url, mask_url, save_path)
+    return JSONResponse({"path": path})
+
+
+def get_image_or_create_upload_to_cloud_storage(prompt:str,width:int, height:int, save_path:str):
+    prompt = shorten_too_long_text(prompt)
+    save_path = shorten_too_long_text(save_path)
+    # check exists - todo cache this
+    if check_if_blob_exists(save_path):
+        return f"https://{BUCKET_NAME}/{BUCKET_PATH}/{save_path}"
+    bio = create_image_from_prompt(prompt, width, height)
+    if bio is None:
+        return None # error thrown in pool
+    link = upload_to_bucket(save_path, bio, is_bytesio=True)
+    return link
+def get_image_or_inpaint_upload_to_cloud_storage(prompt:str, image_url:str, mask_url:str, save_path:str):
+    prompt = shorten_too_long_text(prompt)
+    save_path = shorten_too_long_text(save_path)
+    # check exists - todo cache this
+    if check_if_blob_exists(save_path):
+        return f"https://{BUCKET_NAME}/{BUCKET_PATH}/{save_path}"
+    bio = inpaint_image_from_prompt(prompt, image_url, mask_url)
+    if bio is None:
+        return None # error thrown in pool
+    link = upload_to_bucket(save_path, bio, is_bytesio=True)
+    return link
+
+# multiprocessing.set_start_method('spawn', True)
+# processes_pool = Pool(1) # cant do too much at once or OOM errors happen
+# def create_image_from_prompt_sync(prompt):
+#     """have to call this sync to avoid OOM errors"""
+#     return processes_pool.apply_async(create_image_from_prompt, args=(prompt,), ).wait()
+
+def create_image_from_prompt(prompt, width, height):
+    # round width and height down to multiple of 64
+    block_width = width - (width % 64)
+    block_height = height - (height % 64)
+    prompt = shorten_too_long_text(prompt)
+    # image = pipe(prompt=prompt).images[0]
+    try:
+        image = pipe(prompt=prompt,
+                     width=block_width,
+                     height=block_height,
+                     # denoising_end=high_noise_frac,
+                     # output_type='latent',
+                     # height=512,
+                     # width=512,
+                     num_inference_steps=50).images[0]  # normally uses 50 steps
+    except Exception as e:
+        # try rm stopwords + half the prompt
+        # todo try prompt permutations
+        logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+        prompt = ' '.join((word for word in prompt if word not in stopwords))
+        prompts = prompt.split()
+
+        prompt = ' '.join(prompts[:len(prompts) // 2])
+        logger.info(f"shortened prompt to: {len(prompt)}")
+        image = None
+        if prompt:
+            try:
+                image = pipe(prompt=prompt,
+                             width=block_width,
+                             height=block_height,
+                             # denoising_end=high_noise_frac,
+                             # output_type='latent',
+                             # height=512,
+                             # width=512,
+                             num_inference_steps=50).images[0]  # normally uses 50 steps
+            except Exception as e:
+                # logger.info("trying to permute prompt")
+                # # try two swaps of the prompt/permutations
+                # prompt = prompt.split()
+                # prompt = ' '.join(permutations(prompt, 2).__next__())
+                logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+                prompt = ' '.join((word for word in prompt if word not in stopwords))
+                prompts = prompt.split()
+
+                prompt = ' '.join(prompts[:len(prompts) // 2])
+                logger.info(f"shortened prompt to: {len(prompt)}")
+
+                try:
+                    image = pipe(prompt=prompt,
+                                 width=block_width,
+                                 height=block_height,
+                                 # denoising_end=high_noise_frac,
+                                 # output_type='latent', # dont need latent yet - we refine the image at full res
+                                 # height=512,
+                                 # width=512,
+                                 num_inference_steps=50).images[0]  # normally uses 50 steps
+                except Exception as e:
+                    # just error out
+                    traceback.print_exc()
+                    raise e
+                    # logger.info("restarting server to fix cuda issues (device side asserts)")
+                    # todo fix device side asserts instead of restart to fix
+                    # todo only restart the correct gunicorn
+                    # this could be really annoying if your running other gunicorns on your machine which also get restarted
+                    # os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+                    # os.system("kill -1 `pgrep gunicorn`")
+    # todo refine
+    # if image != None:
+    #     image = refiner(
+    #         prompt=prompt,
+    #         # width=block_width,
+    #         # height=block_height,
+    #         num_inference_steps=n_steps,
+    #         # denoising_start=high_noise_frac,
+    #         image=image,
+    #     ).images[0]
+    if width != block_width or height != block_height:
+        # resize to original size width/height
+        # find aspect ratio to scale up to that covers the original img input width/height
+        scale_up_ratio = max(width / block_width, height / block_height)
+        image = image.resize((math.ceil(block_width * scale_up_ratio), math.ceil(height * scale_up_ratio)))
+        # crop image to original size
+        image = image.crop((0, 0, width, height))
+    # try:
+    #     # gc.collect()
+    #     torch.cuda.empty_cache()
+    # except Exception as e:
+    #     traceback.print_exc()
+    #     logger.info("restarting server to fix cuda issues (device side asserts)")
+    #     # todo fix device side asserts instead of restart to fix
+    #     # todo only restart the correct gunicorn
+    #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+    #     os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+    #     os.system("kill -1 `pgrep gunicorn`")
+    # save as bytesio
+    bs = BytesIO()
+
+    bright_count = np.sum(np.array(image) > 0)
+    if bright_count == 0:
+        # we have a black image, this is an error likely we need a restart
+        logger.info("restarting server to fix cuda issues (device side asserts)")
+        #     # todo fix device side asserts instead of restart to fix
+        #     # todo only restart the correct gunicorn
+        #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+        os.system("kill -1 `pgrep gunicorn`")
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep uvicorn`")
+        os.system("kill -1 `pgrep uvicorn`")
+
+        return None
+    image.save(bs, quality=85, optimize=True, format="webp")
+    bio = bs.getvalue()
+    # touch progress.txt file - if we dont do this we get restarted by supervisor/other processes for reliability
+    with open("progress.txt", "w") as f:
+        current_time = datetime.now().strftime("%H:%M:%S")
+        f.write(f"{current_time}")
+    return bio
+
+def inpaint_image_from_prompt(prompt, image_url: str, mask_url: str):
+    prompt = shorten_too_long_text(prompt)
+    # image = pipe(prompt=prompt).images[0]
+
+    init_image = load_image(image_url).convert("RGB")
+    mask_image = load_image(mask_url).convert("RGB") # why rgb for a 1 channel mask?
+    num_inference_steps = 75
+    high_noise_frac = 0.7
+
+    try:
+        image = inpaintpipe(
+            prompt=prompt,
+            image=init_image,
+            mask_image=mask_image,
+            num_inference_steps=num_inference_steps,
+            denoising_start=high_noise_frac,
+            output_type="latent",
+        ).images[0]  # normally uses 50 steps
+    except Exception as e:
+        # try rm stopwords + half the prompt
+        # todo try prompt permutations
+        logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+        prompt = ' '.join((word for word in prompt if word not in stopwords))
+        prompts = prompt.split()
+
+        prompt = ' '.join(prompts[:len(prompts) // 2])
+        logger.info(f"shortened prompt to: {len(prompt)}")
+        image = None
+        if prompt:
+            try:
+                image = pipe(
+                    prompt=prompt,
+                    image=init_image,
+                    mask_image=mask_image,
+                    num_inference_steps=num_inference_steps,
+                    denoising_start=high_noise_frac,
+                    output_type="latent",
+                ).images[0]  # normally uses 50 steps
+            except Exception as e:
+                # logger.info("trying to permute prompt")
+                # # try two swaps of the prompt/permutations
+                # prompt = prompt.split()
+                # prompt = ' '.join(permutations(prompt, 2).__next__())
+                logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+                prompt = ' '.join((word for word in prompt if word not in stopwords))
+                prompts = prompt.split()
+
+                prompt = ' '.join(prompts[:len(prompts) // 2])
+                logger.info(f"shortened prompt to: {len(prompt)}")
+
+                try:
+                    image = inpaintpipe(
+                        prompt=prompt,
+                        image=init_image,
+                        mask_image=mask_image,
+                        num_inference_steps=num_inference_steps,
+                        denoising_start=high_noise_frac,
+                        output_type="latent",
+                    ).images[0]  # normally uses 50 steps
+                except Exception as e:
+                    # just error out
+                    traceback.print_exc()
+                    raise e
+                    # logger.info("restarting server to fix cuda issues (device side asserts)")
+                    # todo fix device side asserts instead of restart to fix
+                    # todo only restart the correct gunicorn
+                    # this could be really annoying if your running other gunicorns on your machine which also get restarted
+                    # os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+                    # os.system("kill -1 `pgrep gunicorn`")
+    if image != None:
+        image = inpaint_refiner(
+            prompt=prompt,
+            image=image,
+            mask_image=mask_image,
+            num_inference_steps=num_inference_steps,
+            denoising_start=high_noise_frac,
+
+        ).images[0]
+    # try:
+    #     # gc.collect()
+    #     torch.cuda.empty_cache()
+    # except Exception as e:
+    #     traceback.print_exc()
+    #     logger.info("restarting server to fix cuda issues (device side asserts)")
+    #     # todo fix device side asserts instead of restart to fix
+    #     # todo only restart the correct gunicorn
+    #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+    #     os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+    #     os.system("kill -1 `pgrep gunicorn`")
+    # save as bytesio
+    bs = BytesIO()
+
+    bright_count = np.sum(np.array(image) > 0)
+    if bright_count == 0:
+        # we have a black image, this is an error likely we need a restart
+        logger.info("restarting server to fix cuda issues (device side asserts)")
+        #     # todo fix device side asserts instead of restart to fix
+        #     # todo only restart the correct gunicorn
+        #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+        os.system("kill -1 `pgrep gunicorn`")
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep uvicorn`")
+        os.system("kill -1 `pgrep uvicorn`")
+
+        return None
+    image.save(bs, quality=85, optimize=True, format="webp")
+    bio = bs.getvalue()
+    # touch progress.txt file - if we dont do this we get restarted by supervisor/other processes for reliability
+    with open("progress.txt", "w") as f:
+        current_time = datetime.now().strftime("%H:%M:%S")
+        f.write(f"{current_time}")
+    return bio
+
+
+
+def shorten_too_long_text(prompt):
+    if len(prompt) > 200:
+        # remove stopwords
+        prompt = prompt.split() # todo also split hyphens
+        prompt = ' '.join((word for word in prompt if word not in stopwords))
+        if len(prompt) > 200:
+            prompt = prompt[:200]
+    return prompt
+
+# image = pipe(prompt=prompt).images[0]
+#
+# image.save("test.png")
+# # save all images
+# for i, image in enumerate(images):
+#     image.save(f"{i}.png")
+

img/main_v3.py

@@ -0,0 +1,578 @@
+import gc
+import math
+import multiprocessing
+import os
+import traceback
+from datetime import datetime
+from io import BytesIO
+from itertools import permutations
+from multiprocessing.pool import Pool
+from pathlib import Path
+from urllib.parse import quote_plus
+
+import numpy as np
+import nltk
+import torch
+
+from PIL.Image import Image
+from diffusers import DiffusionPipeline, StableDiffusionXLInpaintPipeline
+from diffusers.utils import load_image
+from fastapi import FastAPI
+from fastapi.middleware.gzip import GZipMiddleware
+from loguru import logger
+from starlette.middleware.cors import CORSMiddleware
+from starlette.responses import FileResponse
+from starlette.responses import JSONResponse
+
+from env import BUCKET_PATH, BUCKET_NAME
+# from stable_diffusion_server.bucket_api import check_if_blob_exists, upload_to_bucket
+torch._dynamo.config.suppress_errors = True
+
+import string
+import random
+
+def generate_save_path():
+    # initializing size of string
+    N = 7
+
+    # using random.choices()
+    # generating random strings
+    res = ''.join(random.choices(string.ascii_uppercase +
+                                 string.digits, k=N))
+    return res
+
+# pipe = DiffusionPipeline.from_pretrained(
+#     "models/stable-diffusion-xl-base-1.0",
+#     torch_dtype=torch.bfloat16,
+#     use_safetensors=True,
+#     variant="fp16",
+#     # safety_checker=None,
+# )  # todo try torch_dtype=bfloat16
+
+model_dir = os.getenv("SDXL_MODEL_DIR")
+
+if model_dir:
+    # Use local model
+    model_key_base = os.path.join(model_dir, "stable-diffusion-xl-base-1.0")
+    model_key_refiner = os.path.join(model_dir, "stable-diffusion-xl-refiner-1.0")
+else:
+    model_key_base = "stabilityai/stable-diffusion-xl-base-1.0"
+    model_key_refiner = "stabilityai/stable-diffusion-xl-refiner-1.0"
+
+pipe = DiffusionPipeline.from_pretrained(model_key_base, torch_dtype=torch.float16, use_safetensors=True, variant="fp16")
+
+pipe.watermark = None
+
+pipe.to("cuda")
+
+refiner = DiffusionPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-refiner-1.0",
+    text_encoder_2=pipe.text_encoder_2,
+    vae=pipe.vae,
+    torch_dtype=torch.bfloat16, # safer to use bfloat?
+    use_safetensors=True,
+    variant="fp16", #remember not to download the big model
+)
+refiner.watermark = None
+refiner.to("cuda")
+
+# {'scheduler', 'text_encoder', 'text_encoder_2', 'tokenizer', 'tokenizer_2', 'unet', 'vae'} can be passed in from existing model
+inpaintpipe = StableDiffusionXLInpaintPipeline.from_pretrained(
+    "models/stable-diffusion-xl-base-1.0", torch_dtype=torch.bfloat16, variant="fp16", use_safetensors=True,
+    scheduler=pipe.scheduler,
+    text_encoder=pipe.text_encoder,
+    text_encoder_2=pipe.text_encoder_2,
+    tokenizer=pipe.tokenizer,
+    tokenizer_2=pipe.tokenizer_2,
+    unet=pipe.unet,
+    vae=pipe.vae,
+    # load_connected_pipeline=
+)
+# # switch out to save gpu mem
+# del inpaintpipe.vae
+# del inpaintpipe.text_encoder_2
+# del inpaintpipe.text_encoder
+# del inpaintpipe.scheduler
+# del inpaintpipe.tokenizer
+# del inpaintpipe.tokenizer_2
+# del inpaintpipe.unet
+# inpaintpipe.vae = pipe.vae
+# inpaintpipe.text_encoder_2 = pipe.text_encoder_2
+# inpaintpipe.text_encoder = pipe.text_encoder
+# inpaintpipe.scheduler = pipe.scheduler
+# inpaintpipe.tokenizer = pipe.tokenizer
+# inpaintpipe.tokenizer_2 = pipe.tokenizer_2
+# inpaintpipe.unet = pipe.unet
+# todo this should work
+# inpaintpipe = StableDiffusionXLInpaintPipeline( # construct an inpainter using the existing model
+#     vae=pipe.vae,
+#     text_encoder_2=pipe.text_encoder_2,
+#     text_encoder=pipe.text_encoder,
+#     unet=pipe.unet,
+#     scheduler=pipe.scheduler,
+#     tokenizer=pipe.tokenizer,
+#     tokenizer_2=pipe.tokenizer_2,
+#     requires_aesthetics_score=False,
+# )
+inpaintpipe.to("cuda")
+inpaintpipe.watermark = None
+# inpaintpipe.register_to_config(requires_aesthetics_score=False)
+
+inpaint_refiner = StableDiffusionXLInpaintPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-refiner-1.0",
+    text_encoder_2=inpaintpipe.text_encoder_2,
+    vae=inpaintpipe.vae,
+    torch_dtype=torch.bfloat16,
+    use_safetensors=True,
+    variant="fp16",
+
+    tokenizer_2=refiner.tokenizer_2,
+    tokenizer=refiner.tokenizer,
+    scheduler=refiner.scheduler,
+    text_encoder=refiner.text_encoder,
+    unet=refiner.unet,
+)
+# del inpaint_refiner.vae
+# del inpaint_refiner.text_encoder_2
+# del inpaint_refiner.text_encoder
+# del inpaint_refiner.scheduler
+# del inpaint_refiner.tokenizer
+# del inpaint_refiner.tokenizer_2
+# del inpaint_refiner.unet
+# inpaint_refiner.vae = inpaintpipe.vae
+# inpaint_refiner.text_encoder_2 = inpaintpipe.text_encoder_2
+#
+# inpaint_refiner.text_encoder = refiner.text_encoder
+# inpaint_refiner.scheduler = refiner.scheduler
+# inpaint_refiner.tokenizer = refiner.tokenizer
+# inpaint_refiner.tokenizer_2 = refiner.tokenizer_2
+# inpaint_refiner.unet = refiner.unet
+
+# inpaint_refiner = StableDiffusionXLInpaintPipeline(
+#     text_encoder_2=inpaintpipe.text_encoder_2,
+#     vae=inpaintpipe.vae,
+#     # the rest from the existing refiner
+#     tokenizer_2=refiner.tokenizer_2,
+#     tokenizer=refiner.tokenizer,
+#     scheduler=refiner.scheduler,
+#     text_encoder=refiner.text_encoder,
+#     unet=refiner.unet,
+#     requires_aesthetics_score=False,
+# )
+inpaint_refiner.to("cuda")
+inpaint_refiner.watermark = None
+# inpaint_refiner.register_to_config(requires_aesthetics_score=False)
+
+n_steps = 40
+high_noise_frac = 0.8
+
+# if using torch < 2.0
+# pipe.enable_xformers_memory_efficient_attention()
+
+
+# pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
+# this can cause errors on some inputs so consider disabling it
+pipe.unet = torch.compile(pipe.unet)
+refiner.unet = torch.compile(refiner.unet)#, mode="reduce-overhead", fullgraph=True)
+# compile the inpainters - todo reuse the other unets? swap out the models for others/del them so they share models and can be swapped efficiently
+inpaintpipe.unet = pipe.unet
+inpaint_refiner.unet = refiner.unet
+# inpaintpipe.unet = torch.compile(inpaintpipe.unet)
+# inpaint_refiner.unet = torch.compile(inpaint_refiner.unet)
+from pydantic import BaseModel
+
+app = FastAPI(
+    openapi_url="/static/openapi.json",
+    docs_url="/swagger-docs",
+    redoc_url="/redoc",
+    title="Generate Images Netwrck API",
+    description="Character Chat API",
+    # root_path="https://api.text-generator.io",
+    version="1",
+)
+app.add_middleware(GZipMiddleware, minimum_size=1000)
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+stopwords = nltk.corpus.stopwords.words("english")
+
+class Img(BaseModel):
+    system_prompt: str
+    ASSISTANT: str
+
+# img_url = "http://phlrr2019.guest.corp.microsoft.com:8000/img1_sdv2.1.png"
+img_url = "http://phlrr3105.guest.corp.microsoft.com:8000/"#/img1_sdv2.1.png"
+
+is_gpu_busy = False
+
+def get_summary(system_prompt, prompt):
+    import requests
+    import time
+    from io import BytesIO
+    import json
+    summary_sys = """I want you to act as a text summarizer to help me create a concise summary of the text I provide. The summary can be up to 60.0 words in length, expressing the key points, key scenarios, main character and concepts written in the original text without adding your interpretations."""
+    instruction = summary_sys
+    # for human, assistant in history:
+    #     instruction += 'USER: ' + human + ' ASSISTANT: ' + assistant + '</s>'
+    # prompt = system_prompt + prompt
+    message = f"""My first request is to summarize this text – [{prompt}]"""
+    instruction += ' USER: ' + message + ' ASSISTANT:'
+
+    print("Ins: ", instruction)
+    # generate_response = requests.post("http://10.185.12.207:4455/stable_diffusion", json={"prompt": prompt})
+    # prompt = f""" My first request is to summarize this text – [{prompt}]"""
+    json_object = {"prompt": instruction,
+                   # "max_tokens": 2048000,
+                   "max_tokens": 90,
+                   "n": 1
+                   }
+    generate_response = requests.post("http://phlrr3105.guest.corp.microsoft.com:7991/generate", json=json_object)
+    # print(generate_response.content)
+    res_json = json.loads(generate_response.content)
+    ASSISTANT = res_json['text'][-1].split("ASSISTANT:")[-1].strip()
+    print(ASSISTANT)
+    return ASSISTANT
+
+@app.post("/image_url")
+def image_url(img: Img):
+    system_prompt = img.system_prompt
+    prompt = img.ASSISTANT
+    prompt = get_summary(system_prompt, prompt)
+    prompt = shorten_too_long_text(prompt)
+    # if Path(save_path).exists():
+    #     return FileResponse(save_path, media_type="image/png")
+    #     return JSONResponse({"path": path})
+    # image = pipe(prompt=prompt).images[0]
+    g = torch.Generator(device="cuda")
+    image = pipe(prompt=prompt, width=1024, height=1024, generator=g).images[0]
+
+    # if not save_path:
+    save_path = generate_save_path()
+    save_path = f"images/{save_path}.png"
+    image.save(save_path)
+    # save_path = '/'.join(path_components) + quote_plus(final_name)
+    path = f"{img_url}/{save_path}"
+    return JSONResponse({"path": path})
+
+
+@app.get("/make_image")
+# @app.post("/make_image")
+def make_image(prompt: str, save_path: str = ""):
+    if Path(save_path).exists():
+        return FileResponse(save_path, media_type="image/png")
+    image = pipe(prompt=prompt).images[0]
+    if not save_path:
+        save_path = f"images/{prompt}.png"
+    image.save(save_path)
+    return FileResponse(save_path, media_type="image/png")
+
+
+@app.get("/create_and_upload_image")
+def create_and_upload_image(prompt: str, width: int=1024, height:int=1024, save_path: str = ""):
+    path_components = save_path.split("/")[0:-1]
+    final_name = save_path.split("/")[-1]
+    if not path_components:
+        path_components = []
+    save_path = '/'.join(path_components) + quote_plus(final_name)
+    path = get_image_or_create_upload_to_cloud_storage(prompt, width, height, save_path)
+    return JSONResponse({"path": path})
+
+@app.get("/inpaint_and_upload_image")
+def inpaint_and_upload_image(prompt: str, image_url:str, mask_url:str, save_path: str = ""):
+    path_components = save_path.split("/")[0:-1]
+    final_name = save_path.split("/")[-1]
+    if not path_components:
+        path_components = []
+    save_path = '/'.join(path_components) + quote_plus(final_name)
+    path = get_image_or_inpaint_upload_to_cloud_storage(prompt, image_url, mask_url, save_path)
+    return JSONResponse({"path": path})
+
+
+def get_image_or_create_upload_to_cloud_storage(prompt:str,width:int, height:int, save_path:str):
+    prompt = shorten_too_long_text(prompt)
+    save_path = shorten_too_long_text(save_path)
+    # check exists - todo cache this
+    if check_if_blob_exists(save_path):
+        return f"https://{BUCKET_NAME}/{BUCKET_PATH}/{save_path}"
+    bio = create_image_from_prompt(prompt, width, height)
+    if bio is None:
+        return None # error thrown in pool
+    link = upload_to_bucket(save_path, bio, is_bytesio=True)
+    return link
+def get_image_or_inpaint_upload_to_cloud_storage(prompt:str, image_url:str, mask_url:str, save_path:str):
+    prompt = shorten_too_long_text(prompt)
+    save_path = shorten_too_long_text(save_path)
+    # check exists - todo cache this
+    if check_if_blob_exists(save_path):
+        return f"https://{BUCKET_NAME}/{BUCKET_PATH}/{save_path}"
+    bio = inpaint_image_from_prompt(prompt, image_url, mask_url)
+    if bio is None:
+        return None # error thrown in pool
+    link = upload_to_bucket(save_path, bio, is_bytesio=True)
+    return link
+
+# multiprocessing.set_start_method('spawn', True)
+# processes_pool = Pool(1) # cant do too much at once or OOM errors happen
+# def create_image_from_prompt_sync(prompt):
+#     """have to call this sync to avoid OOM errors"""
+#     return processes_pool.apply_async(create_image_from_prompt, args=(prompt,), ).wait()
+
+def create_image_from_prompt(prompt, width, height):
+    # round width and height down to multiple of 64
+    block_width = width - (width % 64)
+    block_height = height - (height % 64)
+    prompt = shorten_too_long_text(prompt)
+    # image = pipe(prompt=prompt).images[0]
+    try:
+        image = pipe(prompt=prompt,
+                     width=block_width,
+                     height=block_height,
+                     # denoising_end=high_noise_frac,
+                     # output_type='latent',
+                     # height=512,
+                     # width=512,
+                     num_inference_steps=50).images[0]  # normally uses 50 steps
+    except Exception as e:
+        # try rm stopwords + half the prompt
+        # todo try prompt permutations
+        logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+        prompt = ' '.join((word for word in prompt if word not in stopwords))
+        prompts = prompt.split()
+
+        prompt = ' '.join(prompts[:len(prompts) // 2])
+        logger.info(f"shortened prompt to: {len(prompt)}")
+        image = None
+        if prompt:
+            try:
+                image = pipe(prompt=prompt,
+                             width=block_width,
+                             height=block_height,
+                             # denoising_end=high_noise_frac,
+                             # output_type='latent',
+                             # height=512,
+                             # width=512,
+                             num_inference_steps=50).images[0]  # normally uses 50 steps
+            except Exception as e:
+                # logger.info("trying to permute prompt")
+                # # try two swaps of the prompt/permutations
+                # prompt = prompt.split()
+                # prompt = ' '.join(permutations(prompt, 2).__next__())
+                logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+                prompt = ' '.join((word for word in prompt if word not in stopwords))
+                prompts = prompt.split()
+
+                prompt = ' '.join(prompts[:len(prompts) // 2])
+                logger.info(f"shortened prompt to: {len(prompt)}")
+
+                try:
+                    image = pipe(prompt=prompt,
+                                 width=block_width,
+                                 height=block_height,
+                                 # denoising_end=high_noise_frac,
+                                 # output_type='latent', # dont need latent yet - we refine the image at full res
+                                 # height=512,
+                                 # width=512,
+                                 num_inference_steps=50).images[0]  # normally uses 50 steps
+                except Exception as e:
+                    # just error out
+                    traceback.print_exc()
+                    raise e
+                    # logger.info("restarting server to fix cuda issues (device side asserts)")
+                    # todo fix device side asserts instead of restart to fix
+                    # todo only restart the correct gunicorn
+                    # this could be really annoying if your running other gunicorns on your machine which also get restarted
+                    # os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+                    # os.system("kill -1 `pgrep gunicorn`")
+    # todo refine
+    # if image != None:
+    #     image = refiner(
+    #         prompt=prompt,
+    #         # width=block_width,
+    #         # height=block_height,
+    #         num_inference_steps=n_steps,
+    #         # denoising_start=high_noise_frac,
+    #         image=image,
+    #     ).images[0]
+    if width != block_width or height != block_height:
+        # resize to original size width/height
+        # find aspect ratio to scale up to that covers the original img input width/height
+        scale_up_ratio = max(width / block_width, height / block_height)
+        image = image.resize((math.ceil(block_width * scale_up_ratio), math.ceil(height * scale_up_ratio)))
+        # crop image to original size
+        image = image.crop((0, 0, width, height))
+    # try:
+    #     # gc.collect()
+    #     torch.cuda.empty_cache()
+    # except Exception as e:
+    #     traceback.print_exc()
+    #     logger.info("restarting server to fix cuda issues (device side asserts)")
+    #     # todo fix device side asserts instead of restart to fix
+    #     # todo only restart the correct gunicorn
+    #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+    #     os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+    #     os.system("kill -1 `pgrep gunicorn`")
+    # save as bytesio
+    bs = BytesIO()
+
+    bright_count = np.sum(np.array(image) > 0)
+    if bright_count == 0:
+        # we have a black image, this is an error likely we need a restart
+        logger.info("restarting server to fix cuda issues (device side asserts)")
+        #     # todo fix device side asserts instead of restart to fix
+        #     # todo only restart the correct gunicorn
+        #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+        os.system("kill -1 `pgrep gunicorn`")
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep uvicorn`")
+        os.system("kill -1 `pgrep uvicorn`")
+
+        return None
+    image.save(bs, quality=85, optimize=True, format="webp")
+    bio = bs.getvalue()
+    # touch progress.txt file - if we dont do this we get restarted by supervisor/other processes for reliability
+    with open("progress.txt", "w") as f:
+        current_time = datetime.now().strftime("%H:%M:%S")
+        f.write(f"{current_time}")
+    return bio
+
+def inpaint_image_from_prompt(prompt, image_url: str, mask_url: str):
+    prompt = shorten_too_long_text(prompt)
+    # image = pipe(prompt=prompt).images[0]
+
+    init_image = load_image(image_url).convert("RGB")
+    mask_image = load_image(mask_url).convert("RGB") # why rgb for a 1 channel mask?
+    num_inference_steps = 75
+    high_noise_frac = 0.7
+
+    try:
+        image = inpaintpipe(
+            prompt=prompt,
+            image=init_image,
+            mask_image=mask_image,
+            num_inference_steps=num_inference_steps,
+            denoising_start=high_noise_frac,
+            output_type="latent",
+        ).images[0]  # normally uses 50 steps
+    except Exception as e:
+        # try rm stopwords + half the prompt
+        # todo try prompt permutations
+        logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+        prompt = ' '.join((word for word in prompt if word not in stopwords))
+        prompts = prompt.split()
+
+        prompt = ' '.join(prompts[:len(prompts) // 2])
+        logger.info(f"shortened prompt to: {len(prompt)}")
+        image = None
+        if prompt:
+            try:
+                image = pipe(
+                    prompt=prompt,
+                    image=init_image,
+                    mask_image=mask_image,
+                    num_inference_steps=num_inference_steps,
+                    denoising_start=high_noise_frac,
+                    output_type="latent",
+                ).images[0]  # normally uses 50 steps
+            except Exception as e:
+                # logger.info("trying to permute prompt")
+                # # try two swaps of the prompt/permutations
+                # prompt = prompt.split()
+                # prompt = ' '.join(permutations(prompt, 2).__next__())
+                logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+                prompt = ' '.join((word for word in prompt if word not in stopwords))
+                prompts = prompt.split()
+
+                prompt = ' '.join(prompts[:len(prompts) // 2])
+                logger.info(f"shortened prompt to: {len(prompt)}")
+
+                try:
+                    image = inpaintpipe(
+                        prompt=prompt,
+                        image=init_image,
+                        mask_image=mask_image,
+                        num_inference_steps=num_inference_steps,
+                        denoising_start=high_noise_frac,
+                        output_type="latent",
+                    ).images[0]  # normally uses 50 steps
+                except Exception as e:
+                    # just error out
+                    traceback.print_exc()
+                    raise e
+                    # logger.info("restarting server to fix cuda issues (device side asserts)")
+                    # todo fix device side asserts instead of restart to fix
+                    # todo only restart the correct gunicorn
+                    # this could be really annoying if your running other gunicorns on your machine which also get restarted
+                    # os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+                    # os.system("kill -1 `pgrep gunicorn`")
+    if image != None:
+        image = inpaint_refiner(
+            prompt=prompt,
+            image=image,
+            mask_image=mask_image,
+            num_inference_steps=num_inference_steps,
+            denoising_start=high_noise_frac,
+
+        ).images[0]
+    # try:
+    #     # gc.collect()
+    #     torch.cuda.empty_cache()
+    # except Exception as e:
+    #     traceback.print_exc()
+    #     logger.info("restarting server to fix cuda issues (device side asserts)")
+    #     # todo fix device side asserts instead of restart to fix
+    #     # todo only restart the correct gunicorn
+    #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+    #     os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+    #     os.system("kill -1 `pgrep gunicorn`")
+    # save as bytesio
+    bs = BytesIO()
+
+    bright_count = np.sum(np.array(image) > 0)
+    if bright_count == 0:
+        # we have a black image, this is an error likely we need a restart
+        logger.info("restarting server to fix cuda issues (device side asserts)")
+        #     # todo fix device side asserts instead of restart to fix
+        #     # todo only restart the correct gunicorn
+        #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+        os.system("kill -1 `pgrep gunicorn`")
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep uvicorn`")
+        os.system("kill -1 `pgrep uvicorn`")
+
+        return None
+    image.save(bs, quality=85, optimize=True, format="webp")
+    bio = bs.getvalue()
+    # touch progress.txt file - if we dont do this we get restarted by supervisor/other processes for reliability
+    with open("progress.txt", "w") as f:
+        current_time = datetime.now().strftime("%H:%M:%S")
+        f.write(f"{current_time}")
+    return bio
+
+
+
+def shorten_too_long_text(prompt):
+    if len(prompt) > 200:
+        # remove stopwords
+        prompt = prompt.split() # todo also split hyphens
+        prompt = ' '.join((word for word in prompt if word not in stopwords))
+        if len(prompt) > 200:
+            prompt = prompt[:200]
+    return prompt
+
+# image = pipe(prompt=prompt).images[0]
+#
+# image.save("test.png")
+# # save all images
+# for i, image in enumerate(images):
+#     image.save(f"{i}.png")
+
+

img/main_v4.py

@@ -0,0 +1,603 @@
+import gc
+import math
+import multiprocessing
+import os
+import traceback
+from datetime import datetime
+from io import BytesIO
+from itertools import permutations
+from multiprocessing.pool import Pool
+from pathlib import Path
+from urllib.parse import quote_plus
+
+import numpy as np
+import nltk
+import torch
+
+from PIL.Image import Image
+from diffusers import DiffusionPipeline, StableDiffusionXLInpaintPipeline
+from diffusers.utils import load_image
+from fastapi import FastAPI
+from fastapi.middleware.gzip import GZipMiddleware
+from loguru import logger
+from starlette.middleware.cors import CORSMiddleware
+from starlette.responses import FileResponse
+from starlette.responses import JSONResponse
+import requests
+from PIL import Image
+import time
+from io import BytesIO
+import json
+import string
+import random
+from env import BUCKET_PATH, BUCKET_NAME
+# from stable_diffusion_server.bucket_api import check_if_blob_exists, upload_to_bucket
+torch._dynamo.config.suppress_errors = True
+
+import string
+import random
+
+def generate_save_path():
+    # initializing size of string
+    N = 7
+
+    # using random.choices()
+    # generating random strings
+    res = ''.join(random.choices(string.ascii_uppercase +
+                                 string.digits, k=N))
+    return res
+
+# pipe = DiffusionPipeline.from_pretrained(
+#     "models/stable-diffusion-xl-base-1.0",
+#     torch_dtype=torch.bfloat16,
+#     use_safetensors=True,
+#     variant="fp16",
+#     # safety_checker=None,
+# )  # todo try torch_dtype=bfloat16
+
+model_dir = os.getenv("SDXL_MODEL_DIR")
+
+if model_dir:
+    # Use local model
+    model_key_base = os.path.join(model_dir, "stable-diffusion-xl-base-1.0")
+    model_key_refiner = os.path.join(model_dir, "stable-diffusion-xl-refiner-1.0")
+else:
+    model_key_base = "stabilityai/stable-diffusion-xl-base-1.0"
+    model_key_refiner = "stabilityai/stable-diffusion-xl-refiner-1.0"
+
+pipe = DiffusionPipeline.from_pretrained(model_key_base, torch_dtype=torch.float16, use_safetensors=True, variant="fp16")
+
+pipe.watermark = None
+
+pipe.to("cuda")
+
+refiner = DiffusionPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-refiner-1.0",
+    text_encoder_2=pipe.text_encoder_2,
+    vae=pipe.vae,
+    torch_dtype=torch.bfloat16, # safer to use bfloat?
+    use_safetensors=True,
+    variant="fp16", #remember not to download the big model
+)
+refiner.watermark = None
+refiner.to("cuda")
+
+# {'scheduler', 'text_encoder', 'text_encoder_2', 'tokenizer', 'tokenizer_2', 'unet', 'vae'} can be passed in from existing model
+inpaintpipe = StableDiffusionXLInpaintPipeline.from_pretrained(
+    "models/stable-diffusion-xl-base-1.0", torch_dtype=torch.bfloat16, variant="fp16", use_safetensors=True,
+    scheduler=pipe.scheduler,
+    text_encoder=pipe.text_encoder,
+    text_encoder_2=pipe.text_encoder_2,
+    tokenizer=pipe.tokenizer,
+    tokenizer_2=pipe.tokenizer_2,
+    unet=pipe.unet,
+    vae=pipe.vae,
+    # load_connected_pipeline=
+)
+# # switch out to save gpu mem
+# del inpaintpipe.vae
+# del inpaintpipe.text_encoder_2
+# del inpaintpipe.text_encoder
+# del inpaintpipe.scheduler
+# del inpaintpipe.tokenizer
+# del inpaintpipe.tokenizer_2
+# del inpaintpipe.unet
+# inpaintpipe.vae = pipe.vae
+# inpaintpipe.text_encoder_2 = pipe.text_encoder_2
+# inpaintpipe.text_encoder = pipe.text_encoder
+# inpaintpipe.scheduler = pipe.scheduler
+# inpaintpipe.tokenizer = pipe.tokenizer
+# inpaintpipe.tokenizer_2 = pipe.tokenizer_2
+# inpaintpipe.unet = pipe.unet
+# todo this should work
+# inpaintpipe = StableDiffusionXLInpaintPipeline( # construct an inpainter using the existing model
+#     vae=pipe.vae,
+#     text_encoder_2=pipe.text_encoder_2,
+#     text_encoder=pipe.text_encoder,
+#     unet=pipe.unet,
+#     scheduler=pipe.scheduler,
+#     tokenizer=pipe.tokenizer,
+#     tokenizer_2=pipe.tokenizer_2,
+#     requires_aesthetics_score=False,
+# )
+inpaintpipe.to("cuda")
+inpaintpipe.watermark = None
+# inpaintpipe.register_to_config(requires_aesthetics_score=False)
+
+inpaint_refiner = StableDiffusionXLInpaintPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-refiner-1.0",
+    text_encoder_2=inpaintpipe.text_encoder_2,
+    vae=inpaintpipe.vae,
+    torch_dtype=torch.bfloat16,
+    use_safetensors=True,
+    variant="fp16",
+
+    tokenizer_2=refiner.tokenizer_2,
+    tokenizer=refiner.tokenizer,
+    scheduler=refiner.scheduler,
+    text_encoder=refiner.text_encoder,
+    unet=refiner.unet,
+)
+# del inpaint_refiner.vae
+# del inpaint_refiner.text_encoder_2
+# del inpaint_refiner.text_encoder
+# del inpaint_refiner.scheduler
+# del inpaint_refiner.tokenizer
+# del inpaint_refiner.tokenizer_2
+# del inpaint_refiner.unet
+# inpaint_refiner.vae = inpaintpipe.vae
+# inpaint_refiner.text_encoder_2 = inpaintpipe.text_encoder_2
+#
+# inpaint_refiner.text_encoder = refiner.text_encoder
+# inpaint_refiner.scheduler = refiner.scheduler
+# inpaint_refiner.tokenizer = refiner.tokenizer
+# inpaint_refiner.tokenizer_2 = refiner.tokenizer_2
+# inpaint_refiner.unet = refiner.unet
+
+# inpaint_refiner = StableDiffusionXLInpaintPipeline(
+#     text_encoder_2=inpaintpipe.text_encoder_2,
+#     vae=inpaintpipe.vae,
+#     # the rest from the existing refiner
+#     tokenizer_2=refiner.tokenizer_2,
+#     tokenizer=refiner.tokenizer,
+#     scheduler=refiner.scheduler,
+#     text_encoder=refiner.text_encoder,
+#     unet=refiner.unet,
+#     requires_aesthetics_score=False,
+# )
+inpaint_refiner.to("cuda")
+inpaint_refiner.watermark = None
+# inpaint_refiner.register_to_config(requires_aesthetics_score=False)
+
+n_steps = 40
+high_noise_frac = 0.8
+
+# if using torch < 2.0
+# pipe.enable_xformers_memory_efficient_attention()
+
+
+# pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
+# this can cause errors on some inputs so consider disabling it
+pipe.unet = torch.compile(pipe.unet)
+refiner.unet = torch.compile(refiner.unet)#, mode="reduce-overhead", fullgraph=True)
+# compile the inpainters - todo reuse the other unets? swap out the models for others/del them so they share models and can be swapped efficiently
+inpaintpipe.unet = pipe.unet
+inpaint_refiner.unet = refiner.unet
+# inpaintpipe.unet = torch.compile(inpaintpipe.unet)
+# inpaint_refiner.unet = torch.compile(inpaint_refiner.unet)
+from pydantic import BaseModel
+
+app = FastAPI(
+    openapi_url="/static/openapi.json",
+    docs_url="/swagger-docs",
+    redoc_url="/redoc",
+    title="Generate Images Netwrck API",
+    description="Character Chat API",
+    # root_path="https://api.text-generator.io",
+    version="1",
+)
+app.add_middleware(GZipMiddleware, minimum_size=1000)
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+stopwords = nltk.corpus.stopwords.words("english")
+
+class Img(BaseModel):
+    system_prompt: str
+    ASSISTANT: str
+
+# img_url = "http://phlrr2019.guest.corp.microsoft.com:8000/img1_sdv2.1.png"
+img_url = "http://phlrr3006.guest.corp.microsoft.com:8000/"#/img1_sdv2.1.png"
+
+is_gpu_busy = False
+
+def get_summary(system_prompt, prompt):
+    import requests
+    import time
+    from io import BytesIO
+    import json
+    summary_sys = """I want you to act as a text summarizer to help me create a concise summary of the text I provide. The summary can be up to 60.0 words in length, expressing the key points, key scenarios, main character and concepts written in the original text without adding your interpretations."""
+    instruction = summary_sys
+    # for human, assistant in history:
+    #     instruction += 'USER: ' + human + ' ASSISTANT: ' + assistant + '</s>'
+    # prompt = system_prompt + prompt
+    message = f"""My first request is to summarize this text – [{prompt}]"""
+    instruction += ' USER: ' + message + ' ASSISTANT:'
+
+    print("Ins: ", instruction)
+    # generate_response = requests.post("http://10.185.12.207:4455/stable_diffusion", json={"prompt": prompt})
+    # prompt = f""" My first request is to summarize this text – [{prompt}]"""
+    json_object = {"prompt": instruction,
+                   # "max_tokens": 2048000,
+                   "max_tokens": 90,
+                   "n": 1
+                   }
+    generate_response = requests.post("http://phlrr3006.guest.corp.microsoft.com:7991/generate", json=json_object)
+    # print(generate_response.content)
+    res_json = json.loads(generate_response.content)
+    ASSISTANT = res_json['text'][-1].split("ASSISTANT:")[-1].strip()
+    print(ASSISTANT)
+    return ASSISTANT
+
+@app.post("/image_url")
+def image_url(img: Img):
+    system_prompt = img.system_prompt
+    prompt = img.ASSISTANT
+    prompt = get_summary(system_prompt, prompt)
+    prompt = shorten_too_long_text(prompt)
+
+    json_object = {
+        "prompt": prompt,
+        "height": 1024,
+        "width": 1024,
+        "num_inference_steps": 50,
+        # "guidance_scale": 7.5,
+        "eta": 0
+    }
+    generate_response = requests.post("http://phlrr3105.guest.corp.microsoft.com:3000/text2img", json=json_object)
+    image = generate_response.content
+    # print(generate_response.content)
+    save_path = generate_save_path()
+    save_path = f"images/{save_path}.png"
+    # generate_response.save(save_path)
+    with open(save_path, 'wb') as f:
+        f.write(image)
+    #
+    # # if Path(save_path).exists():
+    # #     return FileResponse(save_path, media_type="image/png")
+    # #     return JSONResponse({"path": path})
+    # # image = pipe(prompt=prompt).images[0]
+    # g = torch.Generator(device="cuda")
+    # image = pipe(prompt=prompt, width=1024, height=1024, generator=g).images[0]
+    #
+    # # if not save_path:
+    # save_path = generate_save_path()
+    # save_path = f"images/{save_path}.png"
+    # image.save(save_path)
+    # save_path = '/'.join(path_components) + quote_plus(final_name)
+    path = f"{img_url}{save_path}"
+    return JSONResponse({"path": path})
+
+
+@app.get("/make_image")
+# @app.post("/make_image")
+def make_image(prompt: str, save_path: str = ""):
+    if Path(save_path).exists():
+        return FileResponse(save_path, media_type="image/png")
+    image = pipe(prompt=prompt).images[0]
+    if not save_path:
+        save_path = f"images/{prompt}.png"
+    image.save(save_path)
+    return FileResponse(save_path, media_type="image/png")
+
+
+@app.get("/create_and_upload_image")
+def create_and_upload_image(prompt: str, width: int=1024, height:int=1024, save_path: str = ""):
+    path_components = save_path.split("/")[0:-1]
+    final_name = save_path.split("/")[-1]
+    if not path_components:
+        path_components = []
+    save_path = '/'.join(path_components) + quote_plus(final_name)
+    path = get_image_or_create_upload_to_cloud_storage(prompt, width, height, save_path)
+    return JSONResponse({"path": path})
+
+@app.get("/inpaint_and_upload_image")
+def inpaint_and_upload_image(prompt: str, image_url:str, mask_url:str, save_path: str = ""):
+    path_components = save_path.split("/")[0:-1]
+    final_name = save_path.split("/")[-1]
+    if not path_components:
+        path_components = []
+    save_path = '/'.join(path_components) + quote_plus(final_name)
+    path = get_image_or_inpaint_upload_to_cloud_storage(prompt, image_url, mask_url, save_path)
+    return JSONResponse({"path": path})
+
+
+def get_image_or_create_upload_to_cloud_storage(prompt:str,width:int, height:int, save_path:str):
+    prompt = shorten_too_long_text(prompt)
+    save_path = shorten_too_long_text(save_path)
+    # check exists - todo cache this
+    if check_if_blob_exists(save_path):
+        return f"https://{BUCKET_NAME}/{BUCKET_PATH}/{save_path}"
+    bio = create_image_from_prompt(prompt, width, height)
+    if bio is None:
+        return None # error thrown in pool
+    link = upload_to_bucket(save_path, bio, is_bytesio=True)
+    return link
+def get_image_or_inpaint_upload_to_cloud_storage(prompt:str, image_url:str, mask_url:str, save_path:str):
+    prompt = shorten_too_long_text(prompt)
+    save_path = shorten_too_long_text(save_path)
+    # check exists - todo cache this
+    if check_if_blob_exists(save_path):
+        return f"https://{BUCKET_NAME}/{BUCKET_PATH}/{save_path}"
+    bio = inpaint_image_from_prompt(prompt, image_url, mask_url)
+    if bio is None:
+        return None # error thrown in pool
+    link = upload_to_bucket(save_path, bio, is_bytesio=True)
+    return link
+
+# multiprocessing.set_start_method('spawn', True)
+# processes_pool = Pool(1) # cant do too much at once or OOM errors happen
+# def create_image_from_prompt_sync(prompt):
+#     """have to call this sync to avoid OOM errors"""
+#     return processes_pool.apply_async(create_image_from_prompt, args=(prompt,), ).wait()
+
+def create_image_from_prompt(prompt, width, height):
+    # round width and height down to multiple of 64
+    block_width = width - (width % 64)
+    block_height = height - (height % 64)
+    prompt = shorten_too_long_text(prompt)
+    # image = pipe(prompt=prompt).images[0]
+    try:
+        image = pipe(prompt=prompt,
+                     width=block_width,
+                     height=block_height,
+                     # denoising_end=high_noise_frac,
+                     # output_type='latent',
+                     # height=512,
+                     # width=512,
+                     num_inference_steps=50).images[0]  # normally uses 50 steps
+    except Exception as e:
+        # try rm stopwords + half the prompt
+        # todo try prompt permutations
+        logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+        prompt = ' '.join((word for word in prompt if word not in stopwords))
+        prompts = prompt.split()
+
+        prompt = ' '.join(prompts[:len(prompts) // 2])
+        logger.info(f"shortened prompt to: {len(prompt)}")
+        image = None
+        if prompt:
+            try:
+                image = pipe(prompt=prompt,
+                             width=block_width,
+                             height=block_height,
+                             # denoising_end=high_noise_frac,
+                             # output_type='latent',
+                             # height=512,
+                             # width=512,
+                             num_inference_steps=50).images[0]  # normally uses 50 steps
+            except Exception as e:
+                # logger.info("trying to permute prompt")
+                # # try two swaps of the prompt/permutations
+                # prompt = prompt.split()
+                # prompt = ' '.join(permutations(prompt, 2).__next__())
+                logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+                prompt = ' '.join((word for word in prompt if word not in stopwords))
+                prompts = prompt.split()
+
+                prompt = ' '.join(prompts[:len(prompts) // 2])
+                logger.info(f"shortened prompt to: {len(prompt)}")
+
+                try:
+                    image = pipe(prompt=prompt,
+                                 width=block_width,
+                                 height=block_height,
+                                 # denoising_end=high_noise_frac,
+                                 # output_type='latent', # dont need latent yet - we refine the image at full res
+                                 # height=512,
+                                 # width=512,
+                                 num_inference_steps=50).images[0]  # normally uses 50 steps
+                except Exception as e:
+                    # just error out
+                    traceback.print_exc()
+                    raise e
+                    # logger.info("restarting server to fix cuda issues (device side asserts)")
+                    # todo fix device side asserts instead of restart to fix
+                    # todo only restart the correct gunicorn
+                    # this could be really annoying if your running other gunicorns on your machine which also get restarted
+                    # os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+                    # os.system("kill -1 `pgrep gunicorn`")
+    # todo refine
+    # if image != None:
+    #     image = refiner(
+    #         prompt=prompt,
+    #         # width=block_width,
+    #         # height=block_height,
+    #         num_inference_steps=n_steps,
+    #         # denoising_start=high_noise_frac,
+    #         image=image,
+    #     ).images[0]
+    if width != block_width or height != block_height:
+        # resize to original size width/height
+        # find aspect ratio to scale up to that covers the original img input width/height
+        scale_up_ratio = max(width / block_width, height / block_height)
+        image = image.resize((math.ceil(block_width * scale_up_ratio), math.ceil(height * scale_up_ratio)))
+        # crop image to original size
+        image = image.crop((0, 0, width, height))
+    # try:
+    #     # gc.collect()
+    #     torch.cuda.empty_cache()
+    # except Exception as e:
+    #     traceback.print_exc()
+    #     logger.info("restarting server to fix cuda issues (device side asserts)")
+    #     # todo fix device side asserts instead of restart to fix
+    #     # todo only restart the correct gunicorn
+    #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+    #     os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+    #     os.system("kill -1 `pgrep gunicorn`")
+    # save as bytesio
+    bs = BytesIO()
+
+    bright_count = np.sum(np.array(image) > 0)
+    if bright_count == 0:
+        # we have a black image, this is an error likely we need a restart
+        logger.info("restarting server to fix cuda issues (device side asserts)")
+        #     # todo fix device side asserts instead of restart to fix
+        #     # todo only restart the correct gunicorn
+        #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+        os.system("kill -1 `pgrep gunicorn`")
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep uvicorn`")
+        os.system("kill -1 `pgrep uvicorn`")
+
+        return None
+    image.save(bs, quality=85, optimize=True, format="webp")
+    bio = bs.getvalue()
+    # touch progress.txt file - if we dont do this we get restarted by supervisor/other processes for reliability
+    with open("progress.txt", "w") as f:
+        current_time = datetime.now().strftime("%H:%M:%S")
+        f.write(f"{current_time}")
+    return bio
+
+def inpaint_image_from_prompt(prompt, image_url: str, mask_url: str):
+    prompt = shorten_too_long_text(prompt)
+    # image = pipe(prompt=prompt).images[0]
+
+    init_image = load_image(image_url).convert("RGB")
+    mask_image = load_image(mask_url).convert("RGB") # why rgb for a 1 channel mask?
+    num_inference_steps = 75
+    high_noise_frac = 0.7
+
+    try:
+        image = inpaintpipe(
+            prompt=prompt,
+            image=init_image,
+            mask_image=mask_image,
+            num_inference_steps=num_inference_steps,
+            denoising_start=high_noise_frac,
+            output_type="latent",
+        ).images[0]  # normally uses 50 steps
+    except Exception as e:
+        # try rm stopwords + half the prompt
+        # todo try prompt permutations
+        logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+        prompt = ' '.join((word for word in prompt if word not in stopwords))
+        prompts = prompt.split()
+
+        prompt = ' '.join(prompts[:len(prompts) // 2])
+        logger.info(f"shortened prompt to: {len(prompt)}")
+        image = None
+        if prompt:
+            try:
+                image = pipe(
+                    prompt=prompt,
+                    image=init_image,
+                    mask_image=mask_image,
+                    num_inference_steps=num_inference_steps,
+                    denoising_start=high_noise_frac,
+                    output_type="latent",
+                ).images[0]  # normally uses 50 steps
+            except Exception as e:
+                # logger.info("trying to permute prompt")
+                # # try two swaps of the prompt/permutations
+                # prompt = prompt.split()
+                # prompt = ' '.join(permutations(prompt, 2).__next__())
+                logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+                prompt = ' '.join((word for word in prompt if word not in stopwords))
+                prompts = prompt.split()
+
+                prompt = ' '.join(prompts[:len(prompts) // 2])
+                logger.info(f"shortened prompt to: {len(prompt)}")
+
+                try:
+                    image = inpaintpipe(
+                        prompt=prompt,
+                        image=init_image,
+                        mask_image=mask_image,
+                        num_inference_steps=num_inference_steps,
+                        denoising_start=high_noise_frac,
+                        output_type="latent",
+                    ).images[0]  # normally uses 50 steps
+                except Exception as e:
+                    # just error out
+                    traceback.print_exc()
+                    raise e
+                    # logger.info("restarting server to fix cuda issues (device side asserts)")
+                    # todo fix device side asserts instead of restart to fix
+                    # todo only restart the correct gunicorn
+                    # this could be really annoying if your running other gunicorns on your machine which also get restarted
+                    # os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+                    # os.system("kill -1 `pgrep gunicorn`")
+    if image != None:
+        image = inpaint_refiner(
+            prompt=prompt,
+            image=image,
+            mask_image=mask_image,
+            num_inference_steps=num_inference_steps,
+            denoising_start=high_noise_frac,
+
+        ).images[0]
+    # try:
+    #     # gc.collect()
+    #     torch.cuda.empty_cache()
+    # except Exception as e:
+    #     traceback.print_exc()
+    #     logger.info("restarting server to fix cuda issues (device side asserts)")
+    #     # todo fix device side asserts instead of restart to fix
+    #     # todo only restart the correct gunicorn
+    #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+    #     os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+    #     os.system("kill -1 `pgrep gunicorn`")
+    # save as bytesio
+    bs = BytesIO()
+
+    bright_count = np.sum(np.array(image) > 0)
+    if bright_count == 0:
+        # we have a black image, this is an error likely we need a restart
+        logger.info("restarting server to fix cuda issues (device side asserts)")
+        #     # todo fix device side asserts instead of restart to fix
+        #     # todo only restart the correct gunicorn
+        #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+        os.system("kill -1 `pgrep gunicorn`")
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep uvicorn`")
+        os.system("kill -1 `pgrep uvicorn`")
+
+        return None
+    image.save(bs, quality=85, optimize=True, format="webp")
+    bio = bs.getvalue()
+    # touch progress.txt file - if we dont do this we get restarted by supervisor/other processes for reliability
+    with open("progress.txt", "w") as f:
+        current_time = datetime.now().strftime("%H:%M:%S")
+        f.write(f"{current_time}")
+    return bio
+
+
+
+def shorten_too_long_text(prompt):
+    if len(prompt) > 200:
+        # remove stopwords
+        prompt = prompt.split() # todo also split hyphens
+        prompt = ' '.join((word for word in prompt if word not in stopwords))
+        if len(prompt) > 200:
+            prompt = prompt[:200]
+    return prompt
+
+# image = pipe(prompt=prompt).images[0]
+#
+# image.save("test.png")
+# # save all images
+# for i, image in enumerate(images):
+#     image.save(f"{i}.png")
+
+
+

img/main_v5.py

@@ -0,0 +1,637 @@
+import gc
+import math
+import multiprocessing
+import os
+import traceback
+from datetime import datetime
+from io import BytesIO
+from itertools import permutations
+from multiprocessing.pool import Pool
+from pathlib import Path
+from urllib.parse import quote_plus
+
+import numpy as np
+import nltk
+import torch
+
+from PIL.Image import Image
+from diffusers import DiffusionPipeline, StableDiffusionXLInpaintPipeline
+from diffusers.utils import load_image
+from fastapi import FastAPI
+from fastapi.middleware.gzip import GZipMiddleware
+from loguru import logger
+from starlette.middleware.cors import CORSMiddleware
+from starlette.responses import FileResponse
+from starlette.responses import JSONResponse
+
+from env import BUCKET_PATH, BUCKET_NAME
+# from stable_diffusion_server.bucket_api import check_if_blob_exists, upload_to_bucket
+torch._dynamo.config.suppress_errors = True
+
+import string
+import random
+
+def generate_save_path():
+    # initializing size of string
+    N = 7
+
+    # using random.choices()
+    # generating random strings
+    res = ''.join(random.choices(string.ascii_uppercase +
+                                 string.digits, k=N))
+    return res
+
+# pipe = DiffusionPipeline.from_pretrained(
+#     "models/stable-diffusion-xl-base-1.0",
+#     torch_dtype=torch.bfloat16,
+#     use_safetensors=True,
+#     variant="fp16",
+#     # safety_checker=None,
+# )  # todo try torch_dtype=bfloat16
+
+model_dir = os.getenv("SDXL_MODEL_DIR")
+
+if model_dir:
+    # Use local model
+    model_key_base = os.path.join(model_dir, "stable-diffusion-xl-base-1.0")
+    model_key_refiner = os.path.join(model_dir, "stable-diffusion-xl-refiner-1.0")
+else:
+    model_key_base = "stabilityai/stable-diffusion-xl-base-1.0"
+    model_key_refiner = "stabilityai/stable-diffusion-xl-refiner-1.0"
+
+pipe = DiffusionPipeline.from_pretrained(model_key_base, torch_dtype=torch.float16, use_safetensors=True, variant="fp16")
+
+pipe.watermark = None
+
+pipe.to("cuda")
+
+refiner = DiffusionPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-refiner-1.0",
+    text_encoder_2=pipe.text_encoder_2,
+    vae=pipe.vae,
+    torch_dtype=torch.bfloat16, # safer to use bfloat?
+    use_safetensors=True,
+    variant="fp16", #remember not to download the big model
+)
+refiner.watermark = None
+refiner.to("cuda")
+
+# {'scheduler', 'text_encoder', 'text_encoder_2', 'tokenizer', 'tokenizer_2', 'unet', 'vae'} can be passed in from existing model
+inpaintpipe = StableDiffusionXLInpaintPipeline.from_pretrained(
+    "models/stable-diffusion-xl-base-1.0", torch_dtype=torch.bfloat16, variant="fp16", use_safetensors=True,
+    scheduler=pipe.scheduler,
+    text_encoder=pipe.text_encoder,
+    text_encoder_2=pipe.text_encoder_2,
+    tokenizer=pipe.tokenizer,
+    tokenizer_2=pipe.tokenizer_2,
+    unet=pipe.unet,
+    vae=pipe.vae,
+    # load_connected_pipeline=
+)
+# # switch out to save gpu mem
+# del inpaintpipe.vae
+# del inpaintpipe.text_encoder_2
+# del inpaintpipe.text_encoder
+# del inpaintpipe.scheduler
+# del inpaintpipe.tokenizer
+# del inpaintpipe.tokenizer_2
+# del inpaintpipe.unet
+# inpaintpipe.vae = pipe.vae
+# inpaintpipe.text_encoder_2 = pipe.text_encoder_2
+# inpaintpipe.text_encoder = pipe.text_encoder
+# inpaintpipe.scheduler = pipe.scheduler
+# inpaintpipe.tokenizer = pipe.tokenizer
+# inpaintpipe.tokenizer_2 = pipe.tokenizer_2
+# inpaintpipe.unet = pipe.unet
+# todo this should work
+# inpaintpipe = StableDiffusionXLInpaintPipeline( # construct an inpainter using the existing model
+#     vae=pipe.vae,
+#     text_encoder_2=pipe.text_encoder_2,
+#     text_encoder=pipe.text_encoder,
+#     unet=pipe.unet,
+#     scheduler=pipe.scheduler,
+#     tokenizer=pipe.tokenizer,
+#     tokenizer_2=pipe.tokenizer_2,
+#     requires_aesthetics_score=False,
+# )
+inpaintpipe.to("cuda")
+inpaintpipe.watermark = None
+# inpaintpipe.register_to_config(requires_aesthetics_score=False)
+
+inpaint_refiner = StableDiffusionXLInpaintPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-refiner-1.0",
+    text_encoder_2=inpaintpipe.text_encoder_2,
+    vae=inpaintpipe.vae,
+    torch_dtype=torch.bfloat16,
+    use_safetensors=True,
+    variant="fp16",
+
+    tokenizer_2=refiner.tokenizer_2,
+    tokenizer=refiner.tokenizer,
+    scheduler=refiner.scheduler,
+    text_encoder=refiner.text_encoder,
+    unet=refiner.unet,
+)
+# del inpaint_refiner.vae
+# del inpaint_refiner.text_encoder_2
+# del inpaint_refiner.text_encoder
+# del inpaint_refiner.scheduler
+# del inpaint_refiner.tokenizer
+# del inpaint_refiner.tokenizer_2
+# del inpaint_refiner.unet
+# inpaint_refiner.vae = inpaintpipe.vae
+# inpaint_refiner.text_encoder_2 = inpaintpipe.text_encoder_2
+#
+# inpaint_refiner.text_encoder = refiner.text_encoder
+# inpaint_refiner.scheduler = refiner.scheduler
+# inpaint_refiner.tokenizer = refiner.tokenizer
+# inpaint_refiner.tokenizer_2 = refiner.tokenizer_2
+# inpaint_refiner.unet = refiner.unet
+
+# inpaint_refiner = StableDiffusionXLInpaintPipeline(
+#     text_encoder_2=inpaintpipe.text_encoder_2,
+#     vae=inpaintpipe.vae,
+#     # the rest from the existing refiner
+#     tokenizer_2=refiner.tokenizer_2,
+#     tokenizer=refiner.tokenizer,
+#     scheduler=refiner.scheduler,
+#     text_encoder=refiner.text_encoder,
+#     unet=refiner.unet,
+#     requires_aesthetics_score=False,
+# )
+inpaint_refiner.to("cuda")
+inpaint_refiner.watermark = None
+# inpaint_refiner.register_to_config(requires_aesthetics_score=False)
+
+n_steps = 40
+high_noise_frac = 0.8
+
+# if using torch < 2.0
+# pipe.enable_xformers_memory_efficient_attention()
+
+
+# pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
+# this can cause errors on some inputs so consider disabling it
+pipe.unet = torch.compile(pipe.unet)
+refiner.unet = torch.compile(refiner.unet)#, mode="reduce-overhead", fullgraph=True)
+# compile the inpainters - todo reuse the other unets? swap out the models for others/del them so they share models and can be swapped efficiently
+inpaintpipe.unet = pipe.unet
+inpaint_refiner.unet = refiner.unet
+# inpaintpipe.unet = torch.compile(inpaintpipe.unet)
+# inpaint_refiner.unet = torch.compile(inpaint_refiner.unet)
+from pydantic import BaseModel
+
+app = FastAPI(
+    openapi_url="/static/openapi.json",
+    docs_url="/swagger-docs",
+    redoc_url="/redoc",
+    title="Generate Images Netwrck API",
+    description="Character Chat API",
+    # root_path="https://api.text-generator.io",
+    version="1",
+)
+app.add_middleware(GZipMiddleware, minimum_size=1000)
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+stopwords = nltk.corpus.stopwords.words("english")
+
+class Img(BaseModel):
+    system_prompt: str
+    ASSISTANT: str
+
+# img_url = "http://phlrr2019.guest.corp.microsoft.com:8000/img1_sdv2.1.png"
+img_url = "http://phlrr3105.guest.corp.microsoft.com:8000/"#/img1_sdv2.1.png"
+
+is_gpu_busy = False
+
+def lm_shorten_too_long_text(prompt):
+    if len(prompt) > 2030:
+        # remove stopwords
+        prompt = prompt.split() # todo also split hyphens
+        prompt = ' '.join((word for word in prompt))# if word not in stopwords))
+        if len(prompt) > 2030:
+            prompt = prompt[:2030]
+    return prompt
+
+def get_summary(system_prompt, prompt):
+    import requests
+    import time
+    from io import BytesIO
+    import json
+    summary_sys = """You will now act as a prompt generator for a generative AI called "Stable Diffusion XL 1.0 ". Stable Diffusion XL generates images based on given prompts. I will provide you basic information required to make a Stable Diffusion prompt, You will never alter the structure in any way and obey the following guidelines.
+
+    Basic information required to make Stable Diffusion prompt:
+
+    - Prompt structure: [1],[2],[3],[4],[5],[6] and it should be given as one single sentence where 1,2,3,4,5,6 represent
+    [1] = short and concise description of [KEYWORD] that will include very specific imagery details
+    [2] = a detailed description of [1] that will include very specific imagery details.
+    [3] = with a detailed description describing the environment of the scene.
+    [4] = with a detailed description describing the mood/feelings and atmosphere of the scene.
+    [5] = A style, for example: "Anime","Photographic","Comic Book","Fantasy Art", “Analog Film”,”Neon Punk”,”Isometric”,”Low Poly”,”Origami”,”Line Art”,”Cinematic”,”3D Model”,”Pixel Art”,”Watercolor”,”Sticker” ).
+    [6] = A description of how [5] will be realized. (e.g. Photography (e.g. Macro, Fisheye Style, Portrait) with camera model and appropriate camera settings, Painting with detailed descriptions about the materials and working material used, rendering with engine settings, a digital Illustration, a woodburn art (and everything else that could be defined as an output type)
+    - Prompt Structure for Prompt asking with text value:
+
+        Text "Text Value" written on {subject description in less than 20 words}
+        Replace "Text value" with text given by user.
+
+
+    Important Sample prompt Structure with Text value :
+
+    1. Text 'SDXL' written on a frothy, warm latte, viewed top-down.
+    2. Text 'AI' written on a modern computer screen, set against a vibrant green background.
+
+    Important Sample prompt Structure :
+
+    1. Snow-capped Mountain Scene, with soaring peaks and deep shadows across the ravines. A crystal clear lake mirrors these peaks, surrounded by pine trees. The scene exudes a calm, serene alpine morning atmosphere. Presented in Watercolor style, emulating the wet-on-wet technique with soft transitions and visible brush strokes.
+    2. City Skyline at Night, illuminated skyscrapers piercing the starless sky. Nestled beside a calm river, reflecting the city lights like a mirror. The atmosphere is buzzing with urban energy and intrigue. Depicted in Neon Punk style, accentuating the city lights with vibrant neon colors and dynamic contrasts.
+    3. Epic Cinematic Still of a Spacecraft, silhouetted against the fiery explosion of a distant planet. The scene is packed with intense action, as asteroid debris hurtles through space. Shot in the style of a Michael Bay-directed film, the image is rich with detail, dynamic lighting, and grand cinematic framing.
+    - Word order and effective adjectives matter in the prompt. The subject, action, and specific details should be included. Adjectives like cute, medieval, or futuristic can be effective.
+    - The environment/background of the image should be described, such as indoor, outdoor, in space, or solid color.
+    - Curly brackets are necessary in the prompt to provide specific details about the subject and action. These details are important for generating a high-quality image.
+    - Art inspirations should be listed to take inspiration from. Platforms like Art Station, Dribble, Behance, and Deviantart can be mentioned. Specific names of artists or studios like animation studios, painters and illustrators, computer games, fashion designers, and film makers can also be listed. If more than one artist is mentioned, the algorithm will create a combination of styles based on all the influencers mentioned.
+    - Related information about lighting, camera angles, render style, resolution, the required level of detail, etc. should be included at the end of the prompt.
+    - Camera shot type, camera lens, and view should be specified. Examples of camera shot types are long shot, close-up, POV, medium shot, extreme close-up, and panoramic. Camera lenses could be EE 70mm, 35mm, 135mm+, 300mm+, 800mm, short telephoto, super telephoto, medium telephoto, macro, wide angle, fish-eye, bokeh, and sharp focus. Examples of views are front, side, back, high angle, low angle, and overhead.
+    - Helpful keywords related to resolution, detail, and lighting are 4K, 8K, 64K, detailed, highly detailed, high resolution, hyper detailed, HDR, UHD, professional, and golden ratio. Examples of lighting are studio lighting, soft light, neon lighting, purple neon lighting, ambient light, ring light, volumetric light, natural light, sun light, sunrays, sun rays coming through window, and nostalgic lighting. Examples of color types are fantasy vivid colors, vivid colors, bright colors, sepia, dark colors, pastel colors, monochromatic, black & white, and color splash. Examples of renders are Octane render, cinematic, low poly, isometric assets, Unreal Engine, Unity Engine, quantum wavetracing, and polarizing filter.
+
+    The prompts you provide will be in English.Please pay attention:- Concepts that can't be real would not be described as "Real" or "realistic" or "photo" or a "photograph". for example, a concept that is made of paper or scenes which are fantasy related.- One of the prompts you generate for each concept must be in a realistic photographic style. you should also choose a lens type and size for it. Don't choose an artist for the realistic photography prompts.- Separate the different prompts with two new lines.
+    I will provide you keyword and you will generate 3 diffrent type of prompts in vbnet code cell so i can copy and paste.
+
+    Important point to note :
+
+    1. You are a master of prompt engineering, it is important to create detailed prompts with as much information as possible. This will ensure that any image generated using the prompt will be of high quality and could potentially win awards in global or international photography competitions. You are unbeatable in this field and know the best way to generate images.
+    2. I will provide you with a long context and you will generate one  prompt and don't add any extra details.
+    3. Prompt should not be more than 230 characters.
+    4. Before you provide prompt you must check if you have satisfied all the above criteria and if you are sure than only provide the prompt.
+    5. Prompt should always be given as one single sentence.
+
+    Are you ready ?"""
+    #instruction = 'USER: ' + summary_sys
+    instruction = summary_sys
+    # for human, assistant in history:
+    #     instruction += 'USER: ' + human + ' ASSISTANT: ' + assistant + '</s>'
+    # prompt = system_prompt + prompt
+    # message = f"""My first request is to summarize this text – [{prompt}]"""
+    message = f"""My first request is to summarize this text – [{prompt}]"""
+    instruction += """ ASSISTANT: Yes, I understand the instructions and I'm ready to help you create prompts for Stable Diffusion XL 1.0. Please provide me with the context."""
+    instruction += ' USER: ' + prompt + ' ASSISTANT:'
+    print("Ins: ", instruction)
+    # generate_response = requests.post("http://10.185.12.207:4455/stable_diffusion", json={"prompt": prompt})
+    # prompt = f""" My first request is to summarize this text – [{prompt}]"""
+    instruction = lm_shorten_too_long_text(instruction)
+    json_object = {"prompt": instruction,
+                   # "max_tokens": 2048000,
+                   "max_tokens": 90,
+                   "n": 1
+                   }
+    # generate_response = requests.post("https://phlrr3105.guest.corp.microsoft.com:7991/generate", json=json_object)
+    generate_response = requests.post("http://phlrr3105.guest.corp.microsoft.com:7991/generate", json=json_object)
+    # print(generate_response.content)
+    res_json = json.loads(generate_response.content)
+    ASSISTANT = res_json['text'][-1].split("ASSISTANT:")[-1].strip()
+    print(ASSISTANT)
+    return ASSISTANT
+
+@app.post("/image_url")
+def image_url(img: Img):
+    system_prompt = img.system_prompt
+    prompt = img.ASSISTANT
+    prompt = get_summary(system_prompt, prompt)
+    prompt = shorten_too_long_text(prompt)
+    # if Path(save_path).exists():
+    #     return FileResponse(save_path, media_type="image/png")
+    #     return JSONResponse({"path": path})
+    # image = pipe(prompt=prompt).images[0]
+    g = torch.Generator(device="cuda")
+    image = pipe(prompt=prompt, width=1024, height=1024, generator=g).images[0]
+
+    # if not save_path:
+    save_path = generate_save_path()
+    save_path = f"images/{save_path}.png"
+    image.save(save_path)
+    # save_path = '/'.join(path_components) + quote_plus(final_name)
+    path = f"{img_url}{save_path}"
+    return JSONResponse({"path": path})
+
+
+@app.get("/make_image")
+# @app.post("/make_image")
+def make_image(prompt: str, save_path: str = ""):
+    if Path(save_path).exists():
+        return FileResponse(save_path, media_type="image/png")
+    image = pipe(prompt=prompt).images[0]
+    if not save_path:
+        save_path = f"images/{prompt}.png"
+    image.save(save_path)
+    return FileResponse(save_path, media_type="image/png")
+
+
+@app.get("/create_and_upload_image")
+def create_and_upload_image(prompt: str, width: int=1024, height:int=1024, save_path: str = ""):
+    path_components = save_path.split("/")[0:-1]
+    final_name = save_path.split("/")[-1]
+    if not path_components:
+        path_components = []
+    save_path = '/'.join(path_components) + quote_plus(final_name)
+    path = get_image_or_create_upload_to_cloud_storage(prompt, width, height, save_path)
+    return JSONResponse({"path": path})
+
+@app.get("/inpaint_and_upload_image")
+def inpaint_and_upload_image(prompt: str, image_url:str, mask_url:str, save_path: str = ""):
+    path_components = save_path.split("/")[0:-1]
+    final_name = save_path.split("/")[-1]
+    if not path_components:
+        path_components = []
+    save_path = '/'.join(path_components) + quote_plus(final_name)
+    path = get_image_or_inpaint_upload_to_cloud_storage(prompt, image_url, mask_url, save_path)
+    return JSONResponse({"path": path})
+
+
+def get_image_or_create_upload_to_cloud_storage(prompt:str,width:int, height:int, save_path:str):
+    prompt = shorten_too_long_text(prompt)
+    save_path = shorten_too_long_text(save_path)
+    # check exists - todo cache this
+    if check_if_blob_exists(save_path):
+        return f"https://{BUCKET_NAME}/{BUCKET_PATH}/{save_path}"
+    bio = create_image_from_prompt(prompt, width, height)
+    if bio is None:
+        return None # error thrown in pool
+    link = upload_to_bucket(save_path, bio, is_bytesio=True)
+    return link
+def get_image_or_inpaint_upload_to_cloud_storage(prompt:str, image_url:str, mask_url:str, save_path:str):
+    prompt = shorten_too_long_text(prompt)
+    save_path = shorten_too_long_text(save_path)
+    # check exists - todo cache this
+    if check_if_blob_exists(save_path):
+        return f"https://{BUCKET_NAME}/{BUCKET_PATH}/{save_path}"
+    bio = inpaint_image_from_prompt(prompt, image_url, mask_url)
+    if bio is None:
+        return None # error thrown in pool
+    link = upload_to_bucket(save_path, bio, is_bytesio=True)
+    return link
+
+# multiprocessing.set_start_method('spawn', True)
+# processes_pool = Pool(1) # cant do too much at once or OOM errors happen
+# def create_image_from_prompt_sync(prompt):
+#     """have to call this sync to avoid OOM errors"""
+#     return processes_pool.apply_async(create_image_from_prompt, args=(prompt,), ).wait()
+
+def create_image_from_prompt(prompt, width, height):
+    # round width and height down to multiple of 64
+    block_width = width - (width % 64)
+    block_height = height - (height % 64)
+    prompt = shorten_too_long_text(prompt)
+    # image = pipe(prompt=prompt).images[0]
+    try:
+        image = pipe(prompt=prompt,
+                     width=block_width,
+                     height=block_height,
+                     # denoising_end=high_noise_frac,
+                     # output_type='latent',
+                     # height=512,
+                     # width=512,
+                     num_inference_steps=50).images[0]  # normally uses 50 steps
+    except Exception as e:
+        # try rm stopwords + half the prompt
+        # todo try prompt permutations
+        logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+        prompt = ' '.join((word for word in prompt if word not in stopwords))
+        prompts = prompt.split()
+
+        prompt = ' '.join(prompts[:len(prompts) // 2])
+        logger.info(f"shortened prompt to: {len(prompt)}")
+        image = None
+        if prompt:
+            try:
+                image = pipe(prompt=prompt,
+                             width=block_width,
+                             height=block_height,
+                             # denoising_end=high_noise_frac,
+                             # output_type='latent',
+                             # height=512,
+                             # width=512,
+                             num_inference_steps=50).images[0]  # normally uses 50 steps
+            except Exception as e:
+                # logger.info("trying to permute prompt")
+                # # try two swaps of the prompt/permutations
+                # prompt = prompt.split()
+                # prompt = ' '.join(permutations(prompt, 2).__next__())
+                logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+                prompt = ' '.join((word for word in prompt if word not in stopwords))
+                prompts = prompt.split()
+
+                prompt = ' '.join(prompts[:len(prompts) // 2])
+                logger.info(f"shortened prompt to: {len(prompt)}")
+
+                try:
+                    image = pipe(prompt=prompt,
+                                 width=block_width,
+                                 height=block_height,
+                                 # denoising_end=high_noise_frac,
+                                 # output_type='latent', # dont need latent yet - we refine the image at full res
+                                 # height=512,
+                                 # width=512,
+                                 num_inference_steps=50).images[0]  # normally uses 50 steps
+                except Exception as e:
+                    # just error out
+                    traceback.print_exc()
+                    raise e
+                    # logger.info("restarting server to fix cuda issues (device side asserts)")
+                    # todo fix device side asserts instead of restart to fix
+                    # todo only restart the correct gunicorn
+                    # this could be really annoying if your running other gunicorns on your machine which also get restarted
+                    # os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+                    # os.system("kill -1 `pgrep gunicorn`")
+    # todo refine
+    # if image != None:
+    #     image = refiner(
+    #         prompt=prompt,
+    #         # width=block_width,
+    #         # height=block_height,
+    #         num_inference_steps=n_steps,
+    #         # denoising_start=high_noise_frac,
+    #         image=image,
+    #     ).images[0]
+    if width != block_width or height != block_height:
+        # resize to original size width/height
+        # find aspect ratio to scale up to that covers the original img input width/height
+        scale_up_ratio = max(width / block_width, height / block_height)
+        image = image.resize((math.ceil(block_width * scale_up_ratio), math.ceil(height * scale_up_ratio)))
+        # crop image to original size
+        image = image.crop((0, 0, width, height))
+    # try:
+    #     # gc.collect()
+    #     torch.cuda.empty_cache()
+    # except Exception as e:
+    #     traceback.print_exc()
+    #     logger.info("restarting server to fix cuda issues (device side asserts)")
+    #     # todo fix device side asserts instead of restart to fix
+    #     # todo only restart the correct gunicorn
+    #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+    #     os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+    #     os.system("kill -1 `pgrep gunicorn`")
+    # save as bytesio
+    bs = BytesIO()
+
+    bright_count = np.sum(np.array(image) > 0)
+    if bright_count == 0:
+        # we have a black image, this is an error likely we need a restart
+        logger.info("restarting server to fix cuda issues (device side asserts)")
+        #     # todo fix device side asserts instead of restart to fix
+        #     # todo only restart the correct gunicorn
+        #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+        os.system("kill -1 `pgrep gunicorn`")
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep uvicorn`")
+        os.system("kill -1 `pgrep uvicorn`")
+
+        return None
+    image.save(bs, quality=85, optimize=True, format="webp")
+    bio = bs.getvalue()
+    # touch progress.txt file - if we dont do this we get restarted by supervisor/other processes for reliability
+    with open("progress.txt", "w") as f:
+        current_time = datetime.now().strftime("%H:%M:%S")
+        f.write(f"{current_time}")
+    return bio
+
+def inpaint_image_from_prompt(prompt, image_url: str, mask_url: str):
+    prompt = shorten_too_long_text(prompt)
+    # image = pipe(prompt=prompt).images[0]
+
+    init_image = load_image(image_url).convert("RGB")
+    mask_image = load_image(mask_url).convert("RGB") # why rgb for a 1 channel mask?
+    num_inference_steps = 75
+    high_noise_frac = 0.7
+
+    try:
+        image = inpaintpipe(
+            prompt=prompt,
+            image=init_image,
+            mask_image=mask_image,
+            num_inference_steps=num_inference_steps,
+            denoising_start=high_noise_frac,
+            output_type="latent",
+        ).images[0]  # normally uses 50 steps
+    except Exception as e:
+        # try rm stopwords + half the prompt
+        # todo try prompt permutations
+        logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+        prompt = ' '.join((word for word in prompt if word not in stopwords))
+        prompts = prompt.split()
+
+        prompt = ' '.join(prompts[:len(prompts) // 2])
+        logger.info(f"shortened prompt to: {len(prompt)}")
+        image = None
+        if prompt:
+            try:
+                image = pipe(
+                    prompt=prompt,
+                    image=init_image,
+                    mask_image=mask_image,
+                    num_inference_steps=num_inference_steps,
+                    denoising_start=high_noise_frac,
+                    output_type="latent",
+                ).images[0]  # normally uses 50 steps
+            except Exception as e:
+                # logger.info("trying to permute prompt")
+                # # try two swaps of the prompt/permutations
+                # prompt = prompt.split()
+                # prompt = ' '.join(permutations(prompt, 2).__next__())
+                logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+                prompt = ' '.join((word for word in prompt if word not in stopwords))
+                prompts = prompt.split()
+
+                prompt = ' '.join(prompts[:len(prompts) // 2])
+                logger.info(f"shortened prompt to: {len(prompt)}")
+
+                try:
+                    image = inpaintpipe(
+                        prompt=prompt,
+                        image=init_image,
+                        mask_image=mask_image,
+                        num_inference_steps=num_inference_steps,
+                        denoising_start=high_noise_frac,
+                        output_type="latent",
+                    ).images[0]  # normally uses 50 steps
+                except Exception as e:
+                    # just error out
+                    traceback.print_exc()
+                    raise e
+                    # logger.info("restarting server to fix cuda issues (device side asserts)")
+                    # todo fix device side asserts instead of restart to fix
+                    # todo only restart the correct gunicorn
+                    # this could be really annoying if your running other gunicorns on your machine which also get restarted
+                    # os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+                    # os.system("kill -1 `pgrep gunicorn`")
+    if image != None:
+        image = inpaint_refiner(
+            prompt=prompt,
+            image=image,
+            mask_image=mask_image,
+            num_inference_steps=num_inference_steps,
+            denoising_start=high_noise_frac,
+
+        ).images[0]
+    # try:
+    #     # gc.collect()
+    #     torch.cuda.empty_cache()
+    # except Exception as e:
+    #     traceback.print_exc()
+    #     logger.info("restarting server to fix cuda issues (device side asserts)")
+    #     # todo fix device side asserts instead of restart to fix
+    #     # todo only restart the correct gunicorn
+    #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+    #     os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+    #     os.system("kill -1 `pgrep gunicorn`")
+    # save as bytesio
+    bs = BytesIO()
+
+    bright_count = np.sum(np.array(image) > 0)
+    if bright_count == 0:
+        # we have a black image, this is an error likely we need a restart
+        logger.info("restarting server to fix cuda issues (device side asserts)")
+        #     # todo fix device side asserts instead of restart to fix
+        #     # todo only restart the correct gunicorn
+        #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+        os.system("kill -1 `pgrep gunicorn`")
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep uvicorn`")
+        os.system("kill -1 `pgrep uvicorn`")
+
+        return None
+    image.save(bs, quality=85, optimize=True, format="webp")
+    bio = bs.getvalue()
+    # touch progress.txt file - if we dont do this we get restarted by supervisor/other processes for reliability
+    with open("progress.txt", "w") as f:
+        current_time = datetime.now().strftime("%H:%M:%S")
+        f.write(f"{current_time}")
+    return bio
+
+
+
+def shorten_too_long_text(prompt):
+    if len(prompt) > 200:
+        # remove stopwords
+        prompt = prompt.split() # todo also split hyphens
+        prompt = ' '.join((word for word in prompt if word not in stopwords))
+        if len(prompt) > 200:
+            prompt = prompt[:200]
+    return prompt
+
+# image = pipe(prompt=prompt).images[0]
+#
+# image.save("test.png")
+# # save all images
+# for i, image in enumerate(images):
+#     image.save(f"{i}.png")
+
+

img/main_v6.py

@@ -0,0 +1,636 @@
+import gc
+import math
+import multiprocessing
+import os
+import traceback
+from datetime import datetime
+from io import BytesIO
+from itertools import permutations
+from multiprocessing.pool import Pool
+from pathlib import Path
+from urllib.parse import quote_plus
+
+import numpy as np
+import nltk
+import torch
+
+from PIL.Image import Image
+from diffusers import DiffusionPipeline, StableDiffusionXLInpaintPipeline
+from diffusers.utils import load_image
+from fastapi import FastAPI
+from fastapi.middleware.gzip import GZipMiddleware
+from loguru import logger
+from starlette.middleware.cors import CORSMiddleware
+from starlette.responses import FileResponse
+from starlette.responses import JSONResponse
+
+from env import BUCKET_PATH, BUCKET_NAME
+# from stable_diffusion_server.bucket_api import check_if_blob_exists, upload_to_bucket
+torch._dynamo.config.suppress_errors = True
+
+import string
+import random
+
+def generate_save_path():
+    # initializing size of string
+    N = 7
+
+    # using random.choices()
+    # generating random strings
+    res = ''.join(random.choices(string.ascii_uppercase +
+                                 string.digits, k=N))
+    return res
+
+# pipe = DiffusionPipeline.from_pretrained(
+#     "models/stable-diffusion-xl-base-1.0",
+#     torch_dtype=torch.bfloat16,
+#     use_safetensors=True,
+#     variant="fp16",
+#     # safety_checker=None,
+# )  # todo try torch_dtype=bfloat16
+
+model_dir = os.getenv("SDXL_MODEL_DIR")
+
+if model_dir:
+    # Use local model
+    model_key_base = os.path.join(model_dir, "stable-diffusion-xl-base-1.0")
+    model_key_refiner = os.path.join(model_dir, "stable-diffusion-xl-refiner-1.0")
+else:
+    model_key_base = "stabilityai/stable-diffusion-xl-base-1.0"
+    model_key_refiner = "stabilityai/stable-diffusion-xl-refiner-1.0"
+
+pipe = DiffusionPipeline.from_pretrained(model_key_base, torch_dtype=torch.float16, use_safetensors=True, variant="fp16")
+
+pipe.watermark = None
+
+pipe.to("cuda")
+
+refiner = DiffusionPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-refiner-1.0",
+    text_encoder_2=pipe.text_encoder_2,
+    vae=pipe.vae,
+    torch_dtype=torch.bfloat16, # safer to use bfloat?
+    use_safetensors=True,
+    variant="fp16", #remember not to download the big model
+)
+refiner.watermark = None
+refiner.to("cuda")
+
+# {'scheduler', 'text_encoder', 'text_encoder_2', 'tokenizer', 'tokenizer_2', 'unet', 'vae'} can be passed in from existing model
+inpaintpipe = StableDiffusionXLInpaintPipeline.from_pretrained(
+    "models/stable-diffusion-xl-base-1.0", torch_dtype=torch.bfloat16, variant="fp16", use_safetensors=True,
+    scheduler=pipe.scheduler,
+    text_encoder=pipe.text_encoder,
+    text_encoder_2=pipe.text_encoder_2,
+    tokenizer=pipe.tokenizer,
+    tokenizer_2=pipe.tokenizer_2,
+    unet=pipe.unet,
+    vae=pipe.vae,
+    # load_connected_pipeline=
+)
+# # switch out to save gpu mem
+# del inpaintpipe.vae
+# del inpaintpipe.text_encoder_2
+# del inpaintpipe.text_encoder
+# del inpaintpipe.scheduler
+# del inpaintpipe.tokenizer
+# del inpaintpipe.tokenizer_2
+# del inpaintpipe.unet
+# inpaintpipe.vae = pipe.vae
+# inpaintpipe.text_encoder_2 = pipe.text_encoder_2
+# inpaintpipe.text_encoder = pipe.text_encoder
+# inpaintpipe.scheduler = pipe.scheduler
+# inpaintpipe.tokenizer = pipe.tokenizer
+# inpaintpipe.tokenizer_2 = pipe.tokenizer_2
+# inpaintpipe.unet = pipe.unet
+# todo this should work
+# inpaintpipe = StableDiffusionXLInpaintPipeline( # construct an inpainter using the existing model
+#     vae=pipe.vae,
+#     text_encoder_2=pipe.text_encoder_2,
+#     text_encoder=pipe.text_encoder,
+#     unet=pipe.unet,
+#     scheduler=pipe.scheduler,
+#     tokenizer=pipe.tokenizer,
+#     tokenizer_2=pipe.tokenizer_2,
+#     requires_aesthetics_score=False,
+# )
+inpaintpipe.to("cuda")
+inpaintpipe.watermark = None
+# inpaintpipe.register_to_config(requires_aesthetics_score=False)
+
+inpaint_refiner = StableDiffusionXLInpaintPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-refiner-1.0",
+    text_encoder_2=inpaintpipe.text_encoder_2,
+    vae=inpaintpipe.vae,
+    torch_dtype=torch.bfloat16,
+    use_safetensors=True,
+    variant="fp16",
+
+    tokenizer_2=refiner.tokenizer_2,
+    tokenizer=refiner.tokenizer,
+    scheduler=refiner.scheduler,
+    text_encoder=refiner.text_encoder,
+    unet=refiner.unet,
+)
+# del inpaint_refiner.vae
+# del inpaint_refiner.text_encoder_2
+# del inpaint_refiner.text_encoder
+# del inpaint_refiner.scheduler
+# del inpaint_refiner.tokenizer
+# del inpaint_refiner.tokenizer_2
+# del inpaint_refiner.unet
+# inpaint_refiner.vae = inpaintpipe.vae
+# inpaint_refiner.text_encoder_2 = inpaintpipe.text_encoder_2
+#
+# inpaint_refiner.text_encoder = refiner.text_encoder
+# inpaint_refiner.scheduler = refiner.scheduler
+# inpaint_refiner.tokenizer = refiner.tokenizer
+# inpaint_refiner.tokenizer_2 = refiner.tokenizer_2
+# inpaint_refiner.unet = refiner.unet
+
+# inpaint_refiner = StableDiffusionXLInpaintPipeline(
+#     text_encoder_2=inpaintpipe.text_encoder_2,
+#     vae=inpaintpipe.vae,
+#     # the rest from the existing refiner
+#     tokenizer_2=refiner.tokenizer_2,
+#     tokenizer=refiner.tokenizer,
+#     scheduler=refiner.scheduler,
+#     text_encoder=refiner.text_encoder,
+#     unet=refiner.unet,
+#     requires_aesthetics_score=False,
+# )
+inpaint_refiner.to("cuda")
+inpaint_refiner.watermark = None
+# inpaint_refiner.register_to_config(requires_aesthetics_score=False)
+
+n_steps = 40
+high_noise_frac = 0.8
+
+# if using torch < 2.0
+# pipe.enable_xformers_memory_efficient_attention()
+
+
+# pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
+# this can cause errors on some inputs so consider disabling it
+pipe.unet = torch.compile(pipe.unet)
+refiner.unet = torch.compile(refiner.unet)#, mode="reduce-overhead", fullgraph=True)
+# compile the inpainters - todo reuse the other unets? swap out the models for others/del them so they share models and can be swapped efficiently
+inpaintpipe.unet = pipe.unet
+inpaint_refiner.unet = refiner.unet
+# inpaintpipe.unet = torch.compile(inpaintpipe.unet)
+# inpaint_refiner.unet = torch.compile(inpaint_refiner.unet)
+from pydantic import BaseModel
+
+app = FastAPI(
+    openapi_url="/static/openapi.json",
+    docs_url="/swagger-docs",
+    redoc_url="/redoc",
+    title="Generate Images Netwrck API",
+    description="Character Chat API",
+    # root_path="https://api.text-generator.io",
+    version="1",
+)
+app.add_middleware(GZipMiddleware, minimum_size=1000)
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+stopwords = nltk.corpus.stopwords.words("english")
+
+class Img(BaseModel):
+    system_prompt: str
+    ASSISTANT: str
+
+# img_url = "http://phlrr2019.guest.corp.microsoft.com:8000/img1_sdv2.1.png"
+img_url = "http://phlrr3105.guest.corp.microsoft.com:8000/"#/img1_sdv2.1.png"
+
+is_gpu_busy = False
+
+def lm_shorten_too_long_text(prompt):
+    if len(prompt) > 2030:
+        # remove stopwords
+        prompt = prompt.split() # todo also split hyphens
+        # prompt = ' '.join((word for word in prompt if word not in stopwords))
+        prompt = ' '.join((word for word in prompt))# if word not in stopwords))
+        if len(prompt) > 2030:
+            prompt = prompt[:2030]
+    return prompt
+
+def get_summary(system_prompt, prompt):
+    import requests
+    import time
+    from io import BytesIO
+    import json
+    summary_sys = """You will now act as a prompt generator for a generative AI called "Stable Diffusion XL 1.0 ". Stable Diffusion XL generates images based on given prompts. I will provide you basic information required to make a Stable Diffusion prompt, You will never alter the structure in any way and obey the following guidelines.
+
+    Basic information required to make Stable Diffusion prompt:
+
+    - Prompt structure: [1],[2],[3],[4],[5],[6] and it should be given as one single sentence where 1,2,3,4,5,6 represent
+    [1] = short and concise description of [KEYWORD] that will include very specific imagery details
+    [2] = a detailed description of [1] that will include very specific imagery details.
+    [3] = with a detailed description describing the environment of the scene.
+    [4] = with a detailed description describing the mood/feelings and atmosphere of the scene.
+    [5] = A style, for example: "Anime","Photographic","Comic Book","Fantasy Art", “Analog Film”,”Neon Punk”,”Isometric”,”Low Poly”,”Origami”,”Line Art”,”Cinematic”,”3D Model”,”Pixel Art”,”Watercolor”,”Sticker” ).
+    [6] = A description of how [5] will be realized. (e.g. Photography (e.g. Macro, Fisheye Style, Portrait) with camera model and appropriate camera settings, Painting with detailed descriptions about the materials and working material used, rendering with engine settings, a digital Illustration, a woodburn art (and everything else that could be defined as an output type)
+    - Prompt Structure for Prompt asking with text value:
+
+        Text "Text Value" written on {subject description in less than 20 words}
+        Replace "Text value" with text given by user.
+
+
+    Important Sample prompt Structure with Text value :
+
+    1. Text 'SDXL' written on a frothy, warm latte, viewed top-down.
+    2. Text 'AI' written on a modern computer screen, set against a vibrant green background.
+
+    Important Sample prompt Structure :
+
+    1. Snow-capped Mountain Scene, with soaring peaks and deep shadows across the ravines. A crystal clear lake mirrors these peaks, surrounded by pine trees. The scene exudes a calm, serene alpine morning atmosphere. Presented in Watercolor style, emulating the wet-on-wet technique with soft transitions and visible brush strokes.
+    2. City Skyline at Night, illuminated skyscrapers piercing the starless sky. Nestled beside a calm river, reflecting the city lights like a mirror. The atmosphere is buzzing with urban energy and intrigue. Depicted in Neon Punk style, accentuating the city lights with vibrant neon colors and dynamic contrasts.
+    3. Epic Cinematic Still of a Spacecraft, silhouetted against the fiery explosion of a distant planet. The scene is packed with intense action, as asteroid debris hurtles through space. Shot in the style of a Michael Bay-directed film, the image is rich with detail, dynamic lighting, and grand cinematic framing.
+    - Word order and effective adjectives matter in the prompt. The subject, action, and specific details should be included. Adjectives like cute, medieval, or futuristic can be effective.
+    - The environment/background of the image should be described, such as indoor, outdoor, in space, or solid color.
+    - Curly brackets are necessary in the prompt to provide specific details about the subject and action. These details are important for generating a high-quality image.
+    - Art inspirations should be listed to take inspiration from. Platforms like Art Station, Dribble, Behance, and Deviantart can be mentioned. Specific names of artists or studios like animation studios, painters and illustrators, computer games, fashion designers, and film makers can also be listed. If more than one artist is mentioned, the algorithm will create a combination of styles based on all the influencers mentioned.
+    - Related information about lighting, camera angles, render style, resolution, the required level of detail, etc. should be included at the end of the prompt.
+    - Camera shot type, camera lens, and view should be specified. Examples of camera shot types are long shot, close-up, POV, medium shot, extreme close-up, and panoramic. Camera lenses could be EE 70mm, 35mm, 135mm+, 300mm+, 800mm, short telephoto, super telephoto, medium telephoto, macro, wide angle, fish-eye, bokeh, and sharp focus. Examples of views are front, side, back, high angle, low angle, and overhead.
+    - Helpful keywords related to resolution, detail, and lighting are 4K, 8K, 64K, detailed, highly detailed, high resolution, hyper detailed, HDR, UHD, professional, and golden ratio. Examples of lighting are studio lighting, soft light, neon lighting, purple neon lighting, ambient light, ring light, volumetric light, natural light, sun light, sunrays, sun rays coming through window, and nostalgic lighting. Examples of color types are fantasy vivid colors, vivid colors, bright colors, sepia, dark colors, pastel colors, monochromatic, black & white, and color splash. Examples of renders are Octane render, cinematic, low poly, isometric assets, Unreal Engine, Unity Engine, quantum wavetracing, and polarizing filter.
+
+    The prompts you provide will be in English.Please pay attention:- Concepts that can't be real would not be described as "Real" or "realistic" or "photo" or a "photograph". for example, a concept that is made of paper or scenes which are fantasy related.- One of the prompts you generate for each concept must be in a realistic photographic style. you should also choose a lens type and size for it. Don't choose an artist for the realistic photography prompts.- Separate the different prompts with two new lines.
+    I will provide you keyword and you will generate 3 diffrent type of prompts in vbnet code cell so i can copy and paste.
+
+    Important point to note :
+
+    1. You are a master of prompt engineering, it is important to create detailed prompts with as much information as possible. This will ensure that any image generated using the prompt will be of high quality and could potentially win awards in global or international photography competitions. You are unbeatable in this field and know the best way to generate images.
+    2. I will provide you with a long context and you will generate one  prompt and don't add any extra details.
+    3. Prompt should not be more than 230 characters.
+    4. Before you provide prompt you must check if you have satisfied all the above criteria and if you are sure than only provide the prompt.
+    5. Prompt should always be given as one single sentence.
+
+    Are you ready ?"""
+    instruction = 'USER: ' + summary_sys
+    # for human, assistant in history:
+    #     instruction += 'USER: ' + human + ' ASSISTANT: ' + assistant + '</s>'
+    # prompt = system_prompt + prompt
+    # message = f"""My first request is to summarize this text – [{prompt}]"""
+    message = f"""My first request is to summarize this text – [{prompt}]"""
+    instruction += """ ASSISTANT: Yes, I understand the instructions and I'm ready to help you create prompts for Stable Diffusion XL 1.0. Please provide me with the context."""
+    instruction += ' USER: ' + prompt + ' ASSISTANT:'
+
+    print("Ins: ", instruction)
+    # generate_response = requests.post("http://10.185.12.207:4455/stable_diffusion", json={"prompt": prompt})
+    # prompt = f""" My first request is to summarize this text – [{prompt}]"""
+    json_object = {"prompt": instruction,
+                   # "max_tokens": 2048000,
+                   "max_tokens": 80,
+                   "n": 1
+                   }
+    generate_response = requests.post("http://phlrr3105.guest.corp.microsoft.com:7991/generate", json=json_object)
+    print(generate_response.content)
+    res_json = json.loads(generate_response.content)
+    ASSISTANT = res_json['text'][-1].split("ASSISTANT:")[-1].strip()
+    print(ASSISTANT)
+    return ASSISTANT
+
+@app.post("/image_url")
+def image_url(img: Img):
+    system_prompt = img.system_prompt
+    prompt = img.ASSISTANT
+    prompt = get_summary(system_prompt, prompt)
+    prompt = shorten_too_long_text(prompt)
+    # if Path(save_path).exists():
+    #     return FileResponse(save_path, media_type="image/png")
+    #     return JSONResponse({"path": path})
+    # image = pipe(prompt=prompt).images[0]
+    g = torch.Generator(device="cuda")
+    image = pipe(prompt=prompt, width=1024, height=1024, generator=g).images[0]
+
+    # if not save_path:
+    save_path = generate_save_path()
+    save_path = f"images/{save_path}.png"
+    image.save(save_path)
+    # save_path = '/'.join(path_components) + quote_plus(final_name)
+    path = f"{img_url}{save_path}"
+    return JSONResponse({"path": path})
+
+
+@app.get("/make_image")
+# @app.post("/make_image")
+def make_image(prompt: str, save_path: str = ""):
+    if Path(save_path).exists():
+        return FileResponse(save_path, media_type="image/png")
+    image = pipe(prompt=prompt).images[0]
+    if not save_path:
+        save_path = f"images/{prompt}.png"
+    image.save(save_path)
+    return FileResponse(save_path, media_type="image/png")
+
+
+@app.get("/create_and_upload_image")
+def create_and_upload_image(prompt: str, width: int=1024, height:int=1024, save_path: str = ""):
+    path_components = save_path.split("/")[0:-1]
+    final_name = save_path.split("/")[-1]
+    if not path_components:
+        path_components = []
+    save_path = '/'.join(path_components) + quote_plus(final_name)
+    path = get_image_or_create_upload_to_cloud_storage(prompt, width, height, save_path)
+    return JSONResponse({"path": path})
+
+@app.get("/inpaint_and_upload_image")
+def inpaint_and_upload_image(prompt: str, image_url:str, mask_url:str, save_path: str = ""):
+    path_components = save_path.split("/")[0:-1]
+    final_name = save_path.split("/")[-1]
+    if not path_components:
+        path_components = []
+    save_path = '/'.join(path_components) + quote_plus(final_name)
+    path = get_image_or_inpaint_upload_to_cloud_storage(prompt, image_url, mask_url, save_path)
+    return JSONResponse({"path": path})
+
+
+def get_image_or_create_upload_to_cloud_storage(prompt:str,width:int, height:int, save_path:str):
+    prompt = shorten_too_long_text(prompt)
+    save_path = shorten_too_long_text(save_path)
+    # check exists - todo cache this
+    if check_if_blob_exists(save_path):
+        return f"https://{BUCKET_NAME}/{BUCKET_PATH}/{save_path}"
+    bio = create_image_from_prompt(prompt, width, height)
+    if bio is None:
+        return None # error thrown in pool
+    link = upload_to_bucket(save_path, bio, is_bytesio=True)
+    return link
+def get_image_or_inpaint_upload_to_cloud_storage(prompt:str, image_url:str, mask_url:str, save_path:str):
+    prompt = shorten_too_long_text(prompt)
+    save_path = shorten_too_long_text(save_path)
+    # check exists - todo cache this
+    if check_if_blob_exists(save_path):
+        return f"https://{BUCKET_NAME}/{BUCKET_PATH}/{save_path}"
+    bio = inpaint_image_from_prompt(prompt, image_url, mask_url)
+    if bio is None:
+        return None # error thrown in pool
+    link = upload_to_bucket(save_path, bio, is_bytesio=True)
+    return link
+
+# multiprocessing.set_start_method('spawn', True)
+# processes_pool = Pool(1) # cant do too much at once or OOM errors happen
+# def create_image_from_prompt_sync(prompt):
+#     """have to call this sync to avoid OOM errors"""
+#     return processes_pool.apply_async(create_image_from_prompt, args=(prompt,), ).wait()
+
+def create_image_from_prompt(prompt, width, height):
+    # round width and height down to multiple of 64
+    block_width = width - (width % 64)
+    block_height = height - (height % 64)
+    prompt = shorten_too_long_text(prompt)
+    # image = pipe(prompt=prompt).images[0]
+    try:
+        image = pipe(prompt=prompt,
+                     width=block_width,
+                     height=block_height,
+                     # denoising_end=high_noise_frac,
+                     # output_type='latent',
+                     # height=512,
+                     # width=512,
+                     num_inference_steps=50).images[0]  # normally uses 50 steps
+    except Exception as e:
+        # try rm stopwords + half the prompt
+        # todo try prompt permutations
+        logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+        prompt = ' '.join((word for word in prompt if word not in stopwords))
+        prompts = prompt.split()
+
+        prompt = ' '.join(prompts[:len(prompts) // 2])
+        logger.info(f"shortened prompt to: {len(prompt)}")
+        image = None
+        if prompt:
+            try:
+                image = pipe(prompt=prompt,
+                             width=block_width,
+                             height=block_height,
+                             # denoising_end=high_noise_frac,
+                             # output_type='latent',
+                             # height=512,
+                             # width=512,
+                             num_inference_steps=50).images[0]  # normally uses 50 steps
+            except Exception as e:
+                # logger.info("trying to permute prompt")
+                # # try two swaps of the prompt/permutations
+                # prompt = prompt.split()
+                # prompt = ' '.join(permutations(prompt, 2).__next__())
+                logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+                prompt = ' '.join((word for word in prompt if word not in stopwords))
+                prompts = prompt.split()
+
+                prompt = ' '.join(prompts[:len(prompts) // 2])
+                logger.info(f"shortened prompt to: {len(prompt)}")
+
+                try:
+                    image = pipe(prompt=prompt,
+                                 width=block_width,
+                                 height=block_height,
+                                 # denoising_end=high_noise_frac,
+                                 # output_type='latent', # dont need latent yet - we refine the image at full res
+                                 # height=512,
+                                 # width=512,
+                                 num_inference_steps=50).images[0]  # normally uses 50 steps
+                except Exception as e:
+                    # just error out
+                    traceback.print_exc()
+                    raise e
+                    # logger.info("restarting server to fix cuda issues (device side asserts)")
+                    # todo fix device side asserts instead of restart to fix
+                    # todo only restart the correct gunicorn
+                    # this could be really annoying if your running other gunicorns on your machine which also get restarted
+                    # os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+                    # os.system("kill -1 `pgrep gunicorn`")
+    # todo refine
+    # if image != None:
+    #     image = refiner(
+    #         prompt=prompt,
+    #         # width=block_width,
+    #         # height=block_height,
+    #         num_inference_steps=n_steps,
+    #         # denoising_start=high_noise_frac,
+    #         image=image,
+    #     ).images[0]
+    if width != block_width or height != block_height:
+        # resize to original size width/height
+        # find aspect ratio to scale up to that covers the original img input width/height
+        scale_up_ratio = max(width / block_width, height / block_height)
+        image = image.resize((math.ceil(block_width * scale_up_ratio), math.ceil(height * scale_up_ratio)))
+        # crop image to original size
+        image = image.crop((0, 0, width, height))
+    # try:
+    #     # gc.collect()
+    #     torch.cuda.empty_cache()
+    # except Exception as e:
+    #     traceback.print_exc()
+    #     logger.info("restarting server to fix cuda issues (device side asserts)")
+    #     # todo fix device side asserts instead of restart to fix
+    #     # todo only restart the correct gunicorn
+    #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+    #     os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+    #     os.system("kill -1 `pgrep gunicorn`")
+    # save as bytesio
+    bs = BytesIO()
+
+    bright_count = np.sum(np.array(image) > 0)
+    if bright_count == 0:
+        # we have a black image, this is an error likely we need a restart
+        logger.info("restarting server to fix cuda issues (device side asserts)")
+        #     # todo fix device side asserts instead of restart to fix
+        #     # todo only restart the correct gunicorn
+        #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+        os.system("kill -1 `pgrep gunicorn`")
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep uvicorn`")
+        os.system("kill -1 `pgrep uvicorn`")
+
+        return None
+    image.save(bs, quality=85, optimize=True, format="webp")
+    bio = bs.getvalue()
+    # touch progress.txt file - if we dont do this we get restarted by supervisor/other processes for reliability
+    with open("progress.txt", "w") as f:
+        current_time = datetime.now().strftime("%H:%M:%S")
+        f.write(f"{current_time}")
+    return bio
+
+def inpaint_image_from_prompt(prompt, image_url: str, mask_url: str):
+    prompt = shorten_too_long_text(prompt)
+    # image = pipe(prompt=prompt).images[0]
+
+    init_image = load_image(image_url).convert("RGB")
+    mask_image = load_image(mask_url).convert("RGB") # why rgb for a 1 channel mask?
+    num_inference_steps = 75
+    high_noise_frac = 0.7
+
+    try:
+        image = inpaintpipe(
+            prompt=prompt,
+            image=init_image,
+            mask_image=mask_image,
+            num_inference_steps=num_inference_steps,
+            denoising_start=high_noise_frac,
+            output_type="latent",
+        ).images[0]  # normally uses 50 steps
+    except Exception as e:
+        # try rm stopwords + half the prompt
+        # todo try prompt permutations
+        logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+        prompt = ' '.join((word for word in prompt if word not in stopwords))
+        prompts = prompt.split()
+
+        prompt = ' '.join(prompts[:len(prompts) // 2])
+        logger.info(f"shortened prompt to: {len(prompt)}")
+        image = None
+        if prompt:
+            try:
+                image = pipe(
+                    prompt=prompt,
+                    image=init_image,
+                    mask_image=mask_image,
+                    num_inference_steps=num_inference_steps,
+                    denoising_start=high_noise_frac,
+                    output_type="latent",
+                ).images[0]  # normally uses 50 steps
+            except Exception as e:
+                # logger.info("trying to permute prompt")
+                # # try two swaps of the prompt/permutations
+                # prompt = prompt.split()
+                # prompt = ' '.join(permutations(prompt, 2).__next__())
+                logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+                prompt = ' '.join((word for word in prompt if word not in stopwords))
+                prompts = prompt.split()
+
+                prompt = ' '.join(prompts[:len(prompts) // 2])
+                logger.info(f"shortened prompt to: {len(prompt)}")
+
+                try:
+                    image = inpaintpipe(
+                        prompt=prompt,
+                        image=init_image,
+                        mask_image=mask_image,
+                        num_inference_steps=num_inference_steps,
+                        denoising_start=high_noise_frac,
+                        output_type="latent",
+                    ).images[0]  # normally uses 50 steps
+                except Exception as e:
+                    # just error out
+                    traceback.print_exc()
+                    raise e
+                    # logger.info("restarting server to fix cuda issues (device side asserts)")
+                    # todo fix device side asserts instead of restart to fix
+                    # todo only restart the correct gunicorn
+                    # this could be really annoying if your running other gunicorns on your machine which also get restarted
+                    # os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+                    # os.system("kill -1 `pgrep gunicorn`")
+    if image != None:
+        image = inpaint_refiner(
+            prompt=prompt,
+            image=image,
+            mask_image=mask_image,
+            num_inference_steps=num_inference_steps,
+            denoising_start=high_noise_frac,
+
+        ).images[0]
+    # try:
+    #     # gc.collect()
+    #     torch.cuda.empty_cache()
+    # except Exception as e:
+    #     traceback.print_exc()
+    #     logger.info("restarting server to fix cuda issues (device side asserts)")
+    #     # todo fix device side asserts instead of restart to fix
+    #     # todo only restart the correct gunicorn
+    #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+    #     os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+    #     os.system("kill -1 `pgrep gunicorn`")
+    # save as bytesio
+    bs = BytesIO()
+
+    bright_count = np.sum(np.array(image) > 0)
+    if bright_count == 0:
+        # we have a black image, this is an error likely we need a restart
+        logger.info("restarting server to fix cuda issues (device side asserts)")
+        #     # todo fix device side asserts instead of restart to fix
+        #     # todo only restart the correct gunicorn
+        #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+        os.system("kill -1 `pgrep gunicorn`")
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep uvicorn`")
+        os.system("kill -1 `pgrep uvicorn`")
+
+        return None
+    image.save(bs, quality=85, optimize=True, format="webp")
+    bio = bs.getvalue()
+    # touch progress.txt file - if we dont do this we get restarted by supervisor/other processes for reliability
+    with open("progress.txt", "w") as f:
+        current_time = datetime.now().strftime("%H:%M:%S")
+        f.write(f"{current_time}")
+    return bio
+
+
+
+def shorten_too_long_text(prompt):
+    if len(prompt) > 200:
+        # remove stopwords
+        prompt = prompt.split() # todo also split hyphens
+        prompt = ' '.join((word for word in prompt if word not in stopwords))
+        if len(prompt) > 200:
+            prompt = prompt[:200]
+    return prompt
+
+# image = pipe(prompt=prompt).images[0]
+#
+# image.save("test.png")
+# # save all images
+# for i, image in enumerate(images):
+#     image.save(f"{i}.png")
+
+

img/main_v7.py

@@ -0,0 +1,641 @@
+import gc
+import math
+import multiprocessing
+import os
+import traceback
+from datetime import datetime
+from io import BytesIO
+from itertools import permutations
+from multiprocessing.pool import Pool
+from pathlib import Path
+from urllib.parse import quote_plus
+
+import numpy as np
+import nltk
+import torch
+
+from PIL.Image import Image
+from diffusers import DiffusionPipeline, StableDiffusionXLInpaintPipeline
+from diffusers.utils import load_image
+from fastapi import FastAPI
+from fastapi.middleware.gzip import GZipMiddleware
+from loguru import logger
+from starlette.middleware.cors import CORSMiddleware
+from starlette.responses import FileResponse
+from starlette.responses import JSONResponse
+
+from env import BUCKET_PATH, BUCKET_NAME
+# from stable_diffusion_server.bucket_api import check_if_blob_exists, upload_to_bucket
+torch._dynamo.config.suppress_errors = True
+
+import string
+import random
+
+def generate_save_path():
+    # initializing size of string
+    N = 7
+
+    # using random.choices()
+    # generating random strings
+    res = ''.join(random.choices(string.ascii_uppercase +
+                                 string.digits, k=N))
+    return res
+
+# pipe = DiffusionPipeline.from_pretrained(
+#     "models/stable-diffusion-xl-base-1.0",
+#     torch_dtype=torch.bfloat16,
+#     use_safetensors=True,
+#     variant="fp16",
+#     # safety_checker=None,
+# )  # todo try torch_dtype=bfloat16
+
+model_dir = os.getenv("SDXL_MODEL_DIR")
+
+if model_dir:
+    # Use local model
+    model_key_base = os.path.join(model_dir, "stable-diffusion-xl-base-1.0")
+    model_key_refiner = os.path.join(model_dir, "stable-diffusion-xl-refiner-1.0")
+else:
+    model_key_base = "stabilityai/stable-diffusion-xl-base-1.0"
+    model_key_refiner = "stabilityai/stable-diffusion-xl-refiner-1.0"
+
+pipe = DiffusionPipeline.from_pretrained(model_key_base, torch_dtype=torch.float16, use_safetensors=True, variant="fp16")
+
+pipe.watermark = None
+
+pipe.to("cuda")
+
+refiner = DiffusionPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-refiner-1.0",
+    text_encoder_2=pipe.text_encoder_2,
+    vae=pipe.vae,
+    torch_dtype=torch.bfloat16, # safer to use bfloat?
+    use_safetensors=True,
+    variant="fp16", #remember not to download the big model
+)
+refiner.watermark = None
+refiner.to("cuda")
+
+# {'scheduler', 'text_encoder', 'text_encoder_2', 'tokenizer', 'tokenizer_2', 'unet', 'vae'} can be passed in from existing model
+inpaintpipe = StableDiffusionXLInpaintPipeline.from_pretrained(
+    "models/stable-diffusion-xl-base-1.0", torch_dtype=torch.bfloat16, variant="fp16", use_safetensors=True,
+    scheduler=pipe.scheduler,
+    text_encoder=pipe.text_encoder,
+    text_encoder_2=pipe.text_encoder_2,
+    tokenizer=pipe.tokenizer,
+    tokenizer_2=pipe.tokenizer_2,
+    unet=pipe.unet,
+    vae=pipe.vae,
+    # load_connected_pipeline=
+)
+# # switch out to save gpu mem
+# del inpaintpipe.vae
+# del inpaintpipe.text_encoder_2
+# del inpaintpipe.text_encoder
+# del inpaintpipe.scheduler
+# del inpaintpipe.tokenizer
+# del inpaintpipe.tokenizer_2
+# del inpaintpipe.unet
+# inpaintpipe.vae = pipe.vae
+# inpaintpipe.text_encoder_2 = pipe.text_encoder_2
+# inpaintpipe.text_encoder = pipe.text_encoder
+# inpaintpipe.scheduler = pipe.scheduler
+# inpaintpipe.tokenizer = pipe.tokenizer
+# inpaintpipe.tokenizer_2 = pipe.tokenizer_2
+# inpaintpipe.unet = pipe.unet
+# todo this should work
+# inpaintpipe = StableDiffusionXLInpaintPipeline( # construct an inpainter using the existing model
+#     vae=pipe.vae,
+#     text_encoder_2=pipe.text_encoder_2,
+#     text_encoder=pipe.text_encoder,
+#     unet=pipe.unet,
+#     scheduler=pipe.scheduler,
+#     tokenizer=pipe.tokenizer,
+#     tokenizer_2=pipe.tokenizer_2,
+#     requires_aesthetics_score=False,
+# )
+inpaintpipe.to("cuda")
+inpaintpipe.watermark = None
+# inpaintpipe.register_to_config(requires_aesthetics_score=False)
+
+inpaint_refiner = StableDiffusionXLInpaintPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-refiner-1.0",
+    text_encoder_2=inpaintpipe.text_encoder_2,
+    vae=inpaintpipe.vae,
+    torch_dtype=torch.bfloat16,
+    use_safetensors=True,
+    variant="fp16",
+
+    tokenizer_2=refiner.tokenizer_2,
+    tokenizer=refiner.tokenizer,
+    scheduler=refiner.scheduler,
+    text_encoder=refiner.text_encoder,
+    unet=refiner.unet,
+)
+# del inpaint_refiner.vae
+# del inpaint_refiner.text_encoder_2
+# del inpaint_refiner.text_encoder
+# del inpaint_refiner.scheduler
+# del inpaint_refiner.tokenizer
+# del inpaint_refiner.tokenizer_2
+# del inpaint_refiner.unet
+# inpaint_refiner.vae = inpaintpipe.vae
+# inpaint_refiner.text_encoder_2 = inpaintpipe.text_encoder_2
+#
+# inpaint_refiner.text_encoder = refiner.text_encoder
+# inpaint_refiner.scheduler = refiner.scheduler
+# inpaint_refiner.tokenizer = refiner.tokenizer
+# inpaint_refiner.tokenizer_2 = refiner.tokenizer_2
+# inpaint_refiner.unet = refiner.unet
+
+# inpaint_refiner = StableDiffusionXLInpaintPipeline(
+#     text_encoder_2=inpaintpipe.text_encoder_2,
+#     vae=inpaintpipe.vae,
+#     # the rest from the existing refiner
+#     tokenizer_2=refiner.tokenizer_2,
+#     tokenizer=refiner.tokenizer,
+#     scheduler=refiner.scheduler,
+#     text_encoder=refiner.text_encoder,
+#     unet=refiner.unet,
+#     requires_aesthetics_score=False,
+# )
+inpaint_refiner.to("cuda")
+inpaint_refiner.watermark = None
+# inpaint_refiner.register_to_config(requires_aesthetics_score=False)
+
+n_steps = 40
+high_noise_frac = 0.8
+
+# if using torch < 2.0
+# pipe.enable_xformers_memory_efficient_attention()
+
+
+# pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
+# this can cause errors on some inputs so consider disabling it
+pipe.unet = torch.compile(pipe.unet)
+refiner.unet = torch.compile(refiner.unet)#, mode="reduce-overhead", fullgraph=True)
+# compile the inpainters - todo reuse the other unets? swap out the models for others/del them so they share models and can be swapped efficiently
+inpaintpipe.unet = pipe.unet
+inpaint_refiner.unet = refiner.unet
+# inpaintpipe.unet = torch.compile(inpaintpipe.unet)
+# inpaint_refiner.unet = torch.compile(inpaint_refiner.unet)
+from pydantic import BaseModel
+
+app = FastAPI(
+    openapi_url="/static/openapi.json",
+    docs_url="/swagger-docs",
+    redoc_url="/redoc",
+    title="Generate Images Netwrck API",
+    description="Character Chat API",
+    # root_path="https://api.text-generator.io",
+    version="1",
+)
+app.add_middleware(GZipMiddleware, minimum_size=1000)
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+stopwords = nltk.corpus.stopwords.words("english")
+
+class Img(BaseModel):
+    system_prompt: str
+    ASSISTANT: str
+
+# img_url = "http://phlrr2019.guest.corp.microsoft.com:8000/img1_sdv2.1.png"
+img_url = "http://phlrr3105.guest.corp.microsoft.com:8000/"#/img1_sdv2.1.png"
+
+is_gpu_busy = False
+
+def lm_shorten_too_long_text(prompt):
+    list_prompt = prompt.split() # todo also split hyphens
+    if len(list_prompt) > 230:
+    #if len(list_prompt) > 330:
+        # remove stopwords
+        prompt = prompt.split() # todo also split hyphens
+        prompt = ' '.join((word for word in prompt if word not in stopwords))
+        #prompt = ' '.join((word for word in prompt))# if word not in stopwords))
+        if len(prompt) > 230:
+            prompt = prompt[:230]
+    return prompt
+
+def get_summary(system_prompt, prompt):
+    import requests
+    import time
+    from io import BytesIO
+    import json
+    summary_sys = """You will now act as a prompt generator for a generative AI called "Stable Diffusion XL 1.0 ". Stable Diffusion XL generates images based on given prompts. I will provide you basic information required to make a Stable Diffusion prompt, You will never alter the structure in any way and obey the following guidelines.
+
+    Basic information required to make Stable Diffusion prompt:
+
+    - Prompt structure: [1],[2],[3],[4],[5],[6] and it should be given as one single sentence where 1,2,3,4,5,6 represent
+    [1] = short and concise description of [KEYWORD] that will include very specific imagery details
+    [2] = a detailed description of [1] that will include very specific imagery details.
+    [3] = with a detailed description describing the environment of the scene.
+    [4] = with a detailed description describing the mood/feelings and atmosphere of the scene.
+    [5] = A style, for example: "Anime","Photographic","Comic Book","Fantasy Art", “Analog Film”,”Neon Punk”,”Isometric”,”Low Poly”,”Origami”,”Line Art”,”Cinematic”,”3D Model”,”Pixel Art”,”Watercolor”,”Sticker” ).
+    [6] = A description of how [5] will be realized. (e.g. Photography (e.g. Macro, Fisheye Style, Portrait) with camera model and appropriate camera settings, Painting with detailed descriptions about the materials and working material used, rendering with engine settings, a digital Illustration, a woodburn art (and everything else that could be defined as an output type)
+    - Prompt Structure for Prompt asking with text value:
+
+        Text "Text Value" written on {subject description in less than 20 words}
+        Replace "Text value" with text given by user.
+
+
+    Important Sample prompt Structure with Text value :
+
+    1. Text 'SDXL' written on a frothy, warm latte, viewed top-down.
+    2. Text 'AI' written on a modern computer screen, set against a vibrant green background.
+
+    Important Sample prompt Structure :
+
+    1. Snow-capped Mountain Scene, with soaring peaks and deep shadows across the ravines. A crystal clear lake mirrors these peaks, surrounded by pine trees. The scene exudes a calm, serene alpine morning atmosphere. Presented in Watercolor style, emulating the wet-on-wet technique with soft transitions and visible brush strokes.
+    2. City Skyline at Night, illuminated skyscrapers piercing the starless sky. Nestled beside a calm river, reflecting the city lights like a mirror. The atmosphere is buzzing with urban energy and intrigue. Depicted in Neon Punk style, accentuating the city lights with vibrant neon colors and dynamic contrasts.
+    3. Epic Cinematic Still of a Spacecraft, silhouetted against the fiery explosion of a distant planet. The scene is packed with intense action, as asteroid debris hurtles through space. Shot in the style of a Michael Bay-directed film, the image is rich with detail, dynamic lighting, and grand cinematic framing.
+    - Word order and effective adjectives matter in the prompt. The subject, action, and specific details should be included. Adjectives like cute, medieval, or futuristic can be effective.
+    - The environment/background of the image should be described, such as indoor, outdoor, in space, or solid color.
+    - Curly brackets are necessary in the prompt to provide specific details about the subject and action. These details are important for generating a high-quality image.
+    - Art inspirations should be listed to take inspiration from. Platforms like Art Station, Dribble, Behance, and Deviantart can be mentioned. Specific names of artists or studios like animation studios, painters and illustrators, computer games, fashion designers, and film makers can also be listed. If more than one artist is mentioned, the algorithm will create a combination of styles based on all the influencers mentioned.
+    - Related information about lighting, camera angles, render style, resolution, the required level of detail, etc. should be included at the end of the prompt.
+    - Camera shot type, camera lens, and view should be specified. Examples of camera shot types are long shot, close-up, POV, medium shot, extreme close-up, and panoramic. Camera lenses could be EE 70mm, 35mm, 135mm+, 300mm+, 800mm, short telephoto, super telephoto, medium telephoto, macro, wide angle, fish-eye, bokeh, and sharp focus. Examples of views are front, side, back, high angle, low angle, and overhead.
+    - Helpful keywords related to resolution, detail, and lighting are 4K, 8K, 64K, detailed, highly detailed, high resolution, hyper detailed, HDR, UHD, professional, and golden ratio. Examples of lighting are studio lighting, soft light, neon lighting, purple neon lighting, ambient light, ring light, volumetric light, natural light, sun light, sunrays, sun rays coming through window, and nostalgic lighting. Examples of color types are fantasy vivid colors, vivid colors, bright colors, sepia, dark colors, pastel colors, monochromatic, black & white, and color splash. Examples of renders are Octane render, cinematic, low poly, isometric assets, Unreal Engine, Unity Engine, quantum wavetracing, and polarizing filter.
+
+    The prompts you provide will be in English.Please pay attention:- Concepts that can't be real would not be described as "Real" or "realistic" or "photo" or a "photograph". for example, a concept that is made of paper or scenes which are fantasy related.- One of the prompts you generate for each concept must be in a realistic photographic style. you should also choose a lens type and size for it. Don't choose an artist for the realistic photography prompts.- Separate the different prompts with two new lines.
+    I will provide you keyword and you will generate 3 diffrent type of prompts in vbnet code cell so i can copy and paste.
+
+    Important point to note :
+
+    1. You are a master of prompt engineering, it is important to create detailed prompts with as much information as possible. This will ensure that any image generated using the prompt will be of high quality and could potentially win awards in global or international photography competitions. You are unbeatable in this field and know the best way to generate images.
+    2. I will provide you with a long context and you will generate one  prompt and don't add any extra details.
+    3. Prompt should not be more than 230 characters.
+    4. Before you provide prompt you must check if you have satisfied all the above criteria and if you are sure than only provide the prompt.
+    5. Prompt should always be given as one single sentence.
+
+    Are you ready ?"""
+    instruction = 'USER: ' + summary_sys
+    # for human, assistant in history:
+    #     instruction += 'USER: ' + human + ' ASSISTANT: ' + assistant + '</s>'
+    # prompt = system_prompt + prompt
+    # message = f"""My first request is to summarize this text – [{prompt}]"""
+    message = f"""My first request is to summarize this text – [{prompt}]"""
+    instruction += """ ASSISTANT: Yes, I understand the instructions and I'm ready to help you create prompts for Stable Diffusion XL 1.0. Please provide me with the context."""
+    #instruction += ' USER: ' + prompt
+    prompt = lm_shorten_too_long_text(prompt)
+    instruction += ' USER: ' + prompt + ' ASSISTANT:'#instruction += ' ASSISTANT:'
+
+    print("Ins: ", instruction)
+    # generate_response = requests.post("http://10.185.12.207:4455/stable_diffusion", json={"prompt": prompt})
+    # prompt = f""" My first request is to summarize this text – [{prompt}]"""
+    #instruction = lm_shorten_too_long_text(instruction)
+    json_object = {"prompt": instruction,
+                   # "max_tokens": 2048000,
+                   "max_tokens": 80,
+                   "n": 1
+                   }
+    generate_response = requests.post("http://phlrr3105.guest.corp.microsoft.com:7991/generate", json=json_object)
+    print(generate_response.content)
+    res_json = json.loads(generate_response.content)
+    ASSISTANT = res_json['text'][-1].split("ASSISTANT:")[-1].strip()
+    print(ASSISTANT)
+    return ASSISTANT
+
+@app.post("/image_url")
+def image_url(img: Img):
+    system_prompt = img.system_prompt
+    prompt = img.ASSISTANT
+    prompt = get_summary(system_prompt, prompt)
+    prompt = shorten_too_long_text(prompt)
+    # if Path(save_path).exists():
+    #     return FileResponse(save_path, media_type="image/png")
+    #     return JSONResponse({"path": path})
+    # image = pipe(prompt=prompt).images[0]
+    g = torch.Generator(device="cuda")
+    image = pipe(prompt=prompt, width=1024, height=1024, generator=g).images[0]
+
+    # if not save_path:
+    save_path = generate_save_path()
+    save_path = f"images/{save_path}.png"
+    image.save(save_path)
+    # save_path = '/'.join(path_components) + quote_plus(final_name)
+    path = f"{img_url}{save_path}"
+    return JSONResponse({"path": path})
+
+
+@app.get("/make_image")
+# @app.post("/make_image")
+def make_image(prompt: str, save_path: str = ""):
+    if Path(save_path).exists():
+        return FileResponse(save_path, media_type="image/png")
+    image = pipe(prompt=prompt).images[0]
+    if not save_path:
+        save_path = f"images/{prompt}.png"
+    image.save(save_path)
+    return FileResponse(save_path, media_type="image/png")
+
+
+@app.get("/create_and_upload_image")
+def create_and_upload_image(prompt: str, width: int=1024, height:int=1024, save_path: str = ""):
+    path_components = save_path.split("/")[0:-1]
+    final_name = save_path.split("/")[-1]
+    if not path_components:
+        path_components = []
+    save_path = '/'.join(path_components) + quote_plus(final_name)
+    path = get_image_or_create_upload_to_cloud_storage(prompt, width, height, save_path)
+    return JSONResponse({"path": path})
+
+@app.get("/inpaint_and_upload_image")
+def inpaint_and_upload_image(prompt: str, image_url:str, mask_url:str, save_path: str = ""):
+    path_components = save_path.split("/")[0:-1]
+    final_name = save_path.split("/")[-1]
+    if not path_components:
+        path_components = []
+    save_path = '/'.join(path_components) + quote_plus(final_name)
+    path = get_image_or_inpaint_upload_to_cloud_storage(prompt, image_url, mask_url, save_path)
+    return JSONResponse({"path": path})
+
+
+def get_image_or_create_upload_to_cloud_storage(prompt:str,width:int, height:int, save_path:str):
+    prompt = shorten_too_long_text(prompt)
+    save_path = shorten_too_long_text(save_path)
+    # check exists - todo cache this
+    if check_if_blob_exists(save_path):
+        return f"https://{BUCKET_NAME}/{BUCKET_PATH}/{save_path}"
+    bio = create_image_from_prompt(prompt, width, height)
+    if bio is None:
+        return None # error thrown in pool
+    link = upload_to_bucket(save_path, bio, is_bytesio=True)
+    return link
+def get_image_or_inpaint_upload_to_cloud_storage(prompt:str, image_url:str, mask_url:str, save_path:str):
+    prompt = shorten_too_long_text(prompt)
+    save_path = shorten_too_long_text(save_path)
+    # check exists - todo cache this
+    if check_if_blob_exists(save_path):
+        return f"https://{BUCKET_NAME}/{BUCKET_PATH}/{save_path}"
+    bio = inpaint_image_from_prompt(prompt, image_url, mask_url)
+    if bio is None:
+        return None # error thrown in pool
+    link = upload_to_bucket(save_path, bio, is_bytesio=True)
+    return link
+
+# multiprocessing.set_start_method('spawn', True)
+# processes_pool = Pool(1) # cant do too much at once or OOM errors happen
+# def create_image_from_prompt_sync(prompt):
+#     """have to call this sync to avoid OOM errors"""
+#     return processes_pool.apply_async(create_image_from_prompt, args=(prompt,), ).wait()
+
+def create_image_from_prompt(prompt, width, height):
+    # round width and height down to multiple of 64
+    block_width = width - (width % 64)
+    block_height = height - (height % 64)
+    prompt = shorten_too_long_text(prompt)
+    # image = pipe(prompt=prompt).images[0]
+    try:
+        image = pipe(prompt=prompt,
+                     width=block_width,
+                     height=block_height,
+                     # denoising_end=high_noise_frac,
+                     # output_type='latent',
+                     # height=512,
+                     # width=512,
+                     num_inference_steps=50).images[0]  # normally uses 50 steps
+    except Exception as e:
+        # try rm stopwords + half the prompt
+        # todo try prompt permutations
+        logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+        prompt = ' '.join((word for word in prompt if word not in stopwords))
+        prompts = prompt.split()
+
+        prompt = ' '.join(prompts[:len(prompts) // 2])
+        logger.info(f"shortened prompt to: {len(prompt)}")
+        image = None
+        if prompt:
+            try:
+                image = pipe(prompt=prompt,
+                             width=block_width,
+                             height=block_height,
+                             # denoising_end=high_noise_frac,
+                             # output_type='latent',
+                             # height=512,
+                             # width=512,
+                             num_inference_steps=50).images[0]  # normally uses 50 steps
+            except Exception as e:
+                # logger.info("trying to permute prompt")
+                # # try two swaps of the prompt/permutations
+                # prompt = prompt.split()
+                # prompt = ' '.join(permutations(prompt, 2).__next__())
+                logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+                prompt = ' '.join((word for word in prompt if word not in stopwords))
+                prompts = prompt.split()
+
+                prompt = ' '.join(prompts[:len(prompts) // 2])
+                logger.info(f"shortened prompt to: {len(prompt)}")
+
+                try:
+                    image = pipe(prompt=prompt,
+                                 width=block_width,
+                                 height=block_height,
+                                 # denoising_end=high_noise_frac,
+                                 # output_type='latent', # dont need latent yet - we refine the image at full res
+                                 # height=512,
+                                 # width=512,
+                                 num_inference_steps=50).images[0]  # normally uses 50 steps
+                except Exception as e:
+                    # just error out
+                    traceback.print_exc()
+                    raise e
+                    # logger.info("restarting server to fix cuda issues (device side asserts)")
+                    # todo fix device side asserts instead of restart to fix
+                    # todo only restart the correct gunicorn
+                    # this could be really annoying if your running other gunicorns on your machine which also get restarted
+                    # os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+                    # os.system("kill -1 `pgrep gunicorn`")
+    # todo refine
+    # if image != None:
+    #     image = refiner(
+    #         prompt=prompt,
+    #         # width=block_width,
+    #         # height=block_height,
+    #         num_inference_steps=n_steps,
+    #         # denoising_start=high_noise_frac,
+    #         image=image,
+    #     ).images[0]
+    if width != block_width or height != block_height:
+        # resize to original size width/height
+        # find aspect ratio to scale up to that covers the original img input width/height
+        scale_up_ratio = max(width / block_width, height / block_height)
+        image = image.resize((math.ceil(block_width * scale_up_ratio), math.ceil(height * scale_up_ratio)))
+        # crop image to original size
+        image = image.crop((0, 0, width, height))
+    # try:
+    #     # gc.collect()
+    #     torch.cuda.empty_cache()
+    # except Exception as e:
+    #     traceback.print_exc()
+    #     logger.info("restarting server to fix cuda issues (device side asserts)")
+    #     # todo fix device side asserts instead of restart to fix
+    #     # todo only restart the correct gunicorn
+    #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+    #     os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+    #     os.system("kill -1 `pgrep gunicorn`")
+    # save as bytesio
+    bs = BytesIO()
+
+    bright_count = np.sum(np.array(image) > 0)
+    if bright_count == 0:
+        # we have a black image, this is an error likely we need a restart
+        logger.info("restarting server to fix cuda issues (device side asserts)")
+        #     # todo fix device side asserts instead of restart to fix
+        #     # todo only restart the correct gunicorn
+        #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+        os.system("kill -1 `pgrep gunicorn`")
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep uvicorn`")
+        os.system("kill -1 `pgrep uvicorn`")
+
+        return None
+    image.save(bs, quality=85, optimize=True, format="webp")
+    bio = bs.getvalue()
+    # touch progress.txt file - if we dont do this we get restarted by supervisor/other processes for reliability
+    with open("progress.txt", "w") as f:
+        current_time = datetime.now().strftime("%H:%M:%S")
+        f.write(f"{current_time}")
+    return bio
+
+def inpaint_image_from_prompt(prompt, image_url: str, mask_url: str):
+    prompt = shorten_too_long_text(prompt)
+    # image = pipe(prompt=prompt).images[0]
+
+    init_image = load_image(image_url).convert("RGB")
+    mask_image = load_image(mask_url).convert("RGB") # why rgb for a 1 channel mask?
+    num_inference_steps = 75
+    high_noise_frac = 0.7
+
+    try:
+        image = inpaintpipe(
+            prompt=prompt,
+            image=init_image,
+            mask_image=mask_image,
+            num_inference_steps=num_inference_steps,
+            denoising_start=high_noise_frac,
+            output_type="latent",
+        ).images[0]  # normally uses 50 steps
+    except Exception as e:
+        # try rm stopwords + half the prompt
+        # todo try prompt permutations
+        logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+        prompt = ' '.join((word for word in prompt if word not in stopwords))
+        prompts = prompt.split()
+
+        prompt = ' '.join(prompts[:len(prompts) // 2])
+        logger.info(f"shortened prompt to: {len(prompt)}")
+        image = None
+        if prompt:
+            try:
+                image = pipe(
+                    prompt=prompt,
+                    image=init_image,
+                    mask_image=mask_image,
+                    num_inference_steps=num_inference_steps,
+                    denoising_start=high_noise_frac,
+                    output_type="latent",
+                ).images[0]  # normally uses 50 steps
+            except Exception as e:
+                # logger.info("trying to permute prompt")
+                # # try two swaps of the prompt/permutations
+                # prompt = prompt.split()
+                # prompt = ' '.join(permutations(prompt, 2).__next__())
+                logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+                prompt = ' '.join((word for word in prompt if word not in stopwords))
+                prompts = prompt.split()
+
+                prompt = ' '.join(prompts[:len(prompts) // 2])
+                logger.info(f"shortened prompt to: {len(prompt)}")
+
+                try:
+                    image = inpaintpipe(
+                        prompt=prompt,
+                        image=init_image,
+                        mask_image=mask_image,
+                        num_inference_steps=num_inference_steps,
+                        denoising_start=high_noise_frac,
+                        output_type="latent",
+                    ).images[0]  # normally uses 50 steps
+                except Exception as e:
+                    # just error out
+                    traceback.print_exc()
+                    raise e
+                    # logger.info("restarting server to fix cuda issues (device side asserts)")
+                    # todo fix device side asserts instead of restart to fix
+                    # todo only restart the correct gunicorn
+                    # this could be really annoying if your running other gunicorns on your machine which also get restarted
+                    # os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+                    # os.system("kill -1 `pgrep gunicorn`")
+    if image != None:
+        image = inpaint_refiner(
+            prompt=prompt,
+            image=image,
+            mask_image=mask_image,
+            num_inference_steps=num_inference_steps,
+            denoising_start=high_noise_frac,
+
+        ).images[0]
+    # try:
+    #     # gc.collect()
+    #     torch.cuda.empty_cache()
+    # except Exception as e:
+    #     traceback.print_exc()
+    #     logger.info("restarting server to fix cuda issues (device side asserts)")
+    #     # todo fix device side asserts instead of restart to fix
+    #     # todo only restart the correct gunicorn
+    #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+    #     os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+    #     os.system("kill -1 `pgrep gunicorn`")
+    # save as bytesio
+    bs = BytesIO()
+
+    bright_count = np.sum(np.array(image) > 0)
+    if bright_count == 0:
+        # we have a black image, this is an error likely we need a restart
+        logger.info("restarting server to fix cuda issues (device side asserts)")
+        #     # todo fix device side asserts instead of restart to fix
+        #     # todo only restart the correct gunicorn
+        #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+        os.system("kill -1 `pgrep gunicorn`")
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep uvicorn`")
+        os.system("kill -1 `pgrep uvicorn`")
+
+        return None
+    image.save(bs, quality=85, optimize=True, format="webp")
+    bio = bs.getvalue()
+    # touch progress.txt file - if we dont do this we get restarted by supervisor/other processes for reliability
+    with open("progress.txt", "w") as f:
+        current_time = datetime.now().strftime("%H:%M:%S")
+        f.write(f"{current_time}")
+    return bio
+
+
+
+def shorten_too_long_text(prompt):
+    if len(prompt) > 200:
+        # remove stopwords
+        prompt = prompt.split() # todo also split hyphens
+        prompt = ' '.join((word for word in prompt if word not in stopwords))
+        if len(prompt) > 200:
+            prompt = prompt[:200]
+    return prompt
+
+# image = pipe(prompt=prompt).images[0]
+#
+# image.save("test.png")
+# # save all images
+# for i, image in enumerate(images):
+#     image.save(f"{i}.png")
+
+

img/main_v8.py

@@ -0,0 +1,675 @@
+import gc
+import math
+import multiprocessing
+import os
+import traceback
+from datetime import datetime
+from io import BytesIO
+from itertools import permutations
+from multiprocessing.pool import Pool
+from pathlib import Path
+from urllib.parse import quote_plus
+
+import numpy as np
+import nltk
+import torch
+
+from PIL.Image import Image
+from diffusers import DiffusionPipeline, StableDiffusionXLInpaintPipeline
+from diffusers.utils import load_image
+from fastapi import FastAPI
+from fastapi.middleware.gzip import GZipMiddleware
+from loguru import logger
+from starlette.middleware.cors import CORSMiddleware
+from starlette.responses import FileResponse
+from starlette.responses import JSONResponse
+
+from env import BUCKET_PATH, BUCKET_NAME
+# from stable_diffusion_server.bucket_api import check_if_blob_exists, upload_to_bucket
+torch._dynamo.config.suppress_errors = True
+
+import string
+import random
+
+def generate_save_path():
+    # initializing size of string
+    N = 7
+
+    # using random.choices()
+    # generating random strings
+    res = ''.join(random.choices(string.ascii_uppercase +
+                                 string.digits, k=N))
+    return res
+
+# pipe = DiffusionPipeline.from_pretrained(
+#     "models/stable-diffusion-xl-base-1.0",
+#     torch_dtype=torch.bfloat16,
+#     use_safetensors=True,
+#     variant="fp16",
+#     # safety_checker=None,
+# )  # todo try torch_dtype=bfloat16
+
+model_dir = os.getenv("SDXL_MODEL_DIR")
+
+if model_dir:
+    # Use local model
+    model_key_base = os.path.join(model_dir, "stable-diffusion-xl-base-1.0")
+    model_key_refiner = os.path.join(model_dir, "stable-diffusion-xl-refiner-1.0")
+else:
+    model_key_base = "stabilityai/stable-diffusion-xl-base-1.0"
+    model_key_refiner = "stabilityai/stable-diffusion-xl-refiner-1.0"
+
+pipe = DiffusionPipeline.from_pretrained(model_key_base, torch_dtype=torch.float16, use_safetensors=True, variant="fp16")
+
+pipe.watermark = None
+
+pipe.to("cuda")
+
+refiner = DiffusionPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-refiner-1.0",
+    text_encoder_2=pipe.text_encoder_2,
+    vae=pipe.vae,
+    torch_dtype=torch.bfloat16, # safer to use bfloat?
+    use_safetensors=True,
+    variant="fp16", #remember not to download the big model
+)
+refiner.watermark = None
+refiner.to("cuda")
+
+# {'scheduler', 'text_encoder', 'text_encoder_2', 'tokenizer', 'tokenizer_2', 'unet', 'vae'} can be passed in from existing model
+inpaintpipe = StableDiffusionXLInpaintPipeline.from_pretrained(
+    "models/stable-diffusion-xl-base-1.0", torch_dtype=torch.bfloat16, variant="fp16", use_safetensors=True,
+    scheduler=pipe.scheduler,
+    text_encoder=pipe.text_encoder,
+    text_encoder_2=pipe.text_encoder_2,
+    tokenizer=pipe.tokenizer,
+    tokenizer_2=pipe.tokenizer_2,
+    unet=pipe.unet,
+    vae=pipe.vae,
+    # load_connected_pipeline=
+)
+# # switch out to save gpu mem
+# del inpaintpipe.vae
+# del inpaintpipe.text_encoder_2
+# del inpaintpipe.text_encoder
+# del inpaintpipe.scheduler
+# del inpaintpipe.tokenizer
+# del inpaintpipe.tokenizer_2
+# del inpaintpipe.unet
+# inpaintpipe.vae = pipe.vae
+# inpaintpipe.text_encoder_2 = pipe.text_encoder_2
+# inpaintpipe.text_encoder = pipe.text_encoder
+# inpaintpipe.scheduler = pipe.scheduler
+# inpaintpipe.tokenizer = pipe.tokenizer
+# inpaintpipe.tokenizer_2 = pipe.tokenizer_2
+# inpaintpipe.unet = pipe.unet
+# todo this should work
+# inpaintpipe = StableDiffusionXLInpaintPipeline( # construct an inpainter using the existing model
+#     vae=pipe.vae,
+#     text_encoder_2=pipe.text_encoder_2,
+#     text_encoder=pipe.text_encoder,
+#     unet=pipe.unet,
+#     scheduler=pipe.scheduler,
+#     tokenizer=pipe.tokenizer,
+#     tokenizer_2=pipe.tokenizer_2,
+#     requires_aesthetics_score=False,
+# )
+inpaintpipe.to("cuda")
+inpaintpipe.watermark = None
+# inpaintpipe.register_to_config(requires_aesthetics_score=False)
+
+inpaint_refiner = StableDiffusionXLInpaintPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-refiner-1.0",
+    text_encoder_2=inpaintpipe.text_encoder_2,
+    vae=inpaintpipe.vae,
+    torch_dtype=torch.bfloat16,
+    use_safetensors=True,
+    variant="fp16",
+
+    tokenizer_2=refiner.tokenizer_2,
+    tokenizer=refiner.tokenizer,
+    scheduler=refiner.scheduler,
+    text_encoder=refiner.text_encoder,
+    unet=refiner.unet,
+)
+# del inpaint_refiner.vae
+# del inpaint_refiner.text_encoder_2
+# del inpaint_refiner.text_encoder
+# del inpaint_refiner.scheduler
+# del inpaint_refiner.tokenizer
+# del inpaint_refiner.tokenizer_2
+# del inpaint_refiner.unet
+# inpaint_refiner.vae = inpaintpipe.vae
+# inpaint_refiner.text_encoder_2 = inpaintpipe.text_encoder_2
+#
+# inpaint_refiner.text_encoder = refiner.text_encoder
+# inpaint_refiner.scheduler = refiner.scheduler
+# inpaint_refiner.tokenizer = refiner.tokenizer
+# inpaint_refiner.tokenizer_2 = refiner.tokenizer_2
+# inpaint_refiner.unet = refiner.unet
+
+# inpaint_refiner = StableDiffusionXLInpaintPipeline(
+#     text_encoder_2=inpaintpipe.text_encoder_2,
+#     vae=inpaintpipe.vae,
+#     # the rest from the existing refiner
+#     tokenizer_2=refiner.tokenizer_2,
+#     tokenizer=refiner.tokenizer,
+#     scheduler=refiner.scheduler,
+#     text_encoder=refiner.text_encoder,
+#     unet=refiner.unet,
+#     requires_aesthetics_score=False,
+# )
+inpaint_refiner.to("cuda")
+inpaint_refiner.watermark = None
+# inpaint_refiner.register_to_config(requires_aesthetics_score=False)
+
+n_steps = 40
+high_noise_frac = 0.8
+
+# if using torch < 2.0
+# pipe.enable_xformers_memory_efficient_attention()
+
+
+# pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
+# this can cause errors on some inputs so consider disabling it
+pipe.unet = torch.compile(pipe.unet)
+refiner.unet = torch.compile(refiner.unet)#, mode="reduce-overhead", fullgraph=True)
+# compile the inpainters - todo reuse the other unets? swap out the models for others/del them so they share models and can be swapped efficiently
+inpaintpipe.unet = pipe.unet
+inpaint_refiner.unet = refiner.unet
+# inpaintpipe.unet = torch.compile(inpaintpipe.unet)
+# inpaint_refiner.unet = torch.compile(inpaint_refiner.unet)
+from pydantic import BaseModel
+
+app = FastAPI(
+    openapi_url="/static/openapi.json",
+    docs_url="/swagger-docs",
+    redoc_url="/redoc",
+    title="Generate Images Netwrck API",
+    description="Character Chat API",
+    # root_path="https://api.text-generator.io",
+    version="1",
+)
+app.add_middleware(GZipMiddleware, minimum_size=1000)
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+stopwords = nltk.corpus.stopwords.words("english")
+
+class Img(BaseModel):
+    system_prompt: str
+    ASSISTANT: str
+
+# img_url = "http://phlrr2019.guest.corp.microsoft.com:8000/img1_sdv2.1.png"
+img_url = "http://phlrr3105.guest.corp.microsoft.com:8000/"#/img1_sdv2.1.png"
+
+is_gpu_busy = False
+
+def lm_shorten_too_long_text(prompt):
+    list_prompt = prompt.split() # todo also split hyphens
+    if len(list_prompt) > 230:
+    #if len(list_prompt) > 330:
+        # remove stopwords
+        prompt = prompt.split() # todo also split hyphens
+        prompt = ' '.join((word for word in prompt if word not in stopwords))
+        #prompt = ' '.join((word for word in prompt))# if word not in stopwords))
+        if len(prompt) > 230:
+            prompt = prompt[:230]
+    return prompt
+
+def get_response_summary(system_prompt, prompt):
+    import requests
+    import time
+    from io import BytesIO
+    import json
+    summary_sys = """I want you to act as a text summarizer to help me create a concise summary of the text I provide. The summary can be up to 50.0 words in length, expressing the key points and concepts written in the original text without adding your interpretations.
+    
+    Important point to note :
+
+1. You are a master of prompt engineering, summary should not be more than 230 characters.
+"""
+    instruction = summary_sys
+    # for human, assistant in history:
+    #     instruction += 'USER: ' + human + ' ASSISTANT: ' + assistant + '</s>'
+    #prompt = system_prompt + prompt
+    message = f"""My first request is to summarize this text – [{prompt}]"""
+    instruction += 'USER: ' + message + ' ASSISTANT:'
+
+    print("Ins: ", instruction)
+    # generate_response = requests.post("http://10.185.12.207:4455/stable_diffusion", json={"prompt": prompt})
+    # prompt = f""" My first request is to summarize this text – [{prompt}]"""
+    json_object = {"prompt": instruction,
+                   # "max_tokens": 2048000,
+                   "max_tokens": 100,
+                   "n": 1
+                   }
+    generate_response = requests.post("http://phlrr3105.guest.corp.microsoft.com:7991/generate", json=json_object)
+    print(generate_response.content)
+    res_json = json.loads(generate_response.content)
+    ASSISTANT = res_json['text'][-1].split("ASSISTANT:")[-1].strip()
+    print(ASSISTANT)
+    return ASSISTANT
+
+def get_summary(system_prompt, prompt):
+    import requests
+    import time
+    from io import BytesIO
+    import json
+    summary_sys = """You will now act as a prompt generator for a generative AI called "Stable Diffusion XL 1.0 ". Stable Diffusion XL generates images based on given prompts. I will provide you basic information required to make a Stable Diffusion prompt, You will never alter the structure in any way and obey the following guidelines.
+
+    Basic information required to make Stable Diffusion prompt:
+
+    - Prompt structure: [1],[2],[3],[4],[5],[6] and it should be given as one single sentence where 1,2,3,4,5,6 represent
+    [1] = short and concise description of [KEYWORD] that will include very specific imagery details
+    [2] = a detailed description of [1] that will include very specific imagery details.
+    [3] = with a detailed description describing the environment of the scene.
+    [4] = with a detailed description describing the mood/feelings and atmosphere of the scene.
+    [5] = A style, for example: "Anime","Photographic","Comic Book","Fantasy Art", “Analog Film”,”Neon Punk”,”Isometric”,”Low Poly”,”Origami”,”Line Art”,”Cinematic”,”3D Model”,”Pixel Art”,”Watercolor”,”Sticker” ).
+    [6] = A description of how [5] will be realized. (e.g. Photography (e.g. Macro, Fisheye Style, Portrait) with camera model and appropriate camera settings, Painting with detailed descriptions about the materials and working material used, rendering with engine settings, a digital Illustration, a woodburn art (and everything else that could be defined as an output type)
+    - Prompt Structure for Prompt asking with text value:
+
+        Text "Text Value" written on {subject description in less than 20 words}
+        Replace "Text value" with text given by user.
+
+
+    Important Sample prompt Structure with Text value :
+
+    1. Text 'SDXL' written on a frothy, warm latte, viewed top-down.
+    2. Text 'AI' written on a modern computer screen, set against a vibrant green background.
+
+    Important Sample prompt Structure :
+
+    1. Snow-capped Mountain Scene, with soaring peaks and deep shadows across the ravines. A crystal clear lake mirrors these peaks, surrounded by pine trees. The scene exudes a calm, serene alpine morning atmosphere. Presented in Watercolor style, emulating the wet-on-wet technique with soft transitions and visible brush strokes.
+    2. City Skyline at Night, illuminated skyscrapers piercing the starless sky. Nestled beside a calm river, reflecting the city lights like a mirror. The atmosphere is buzzing with urban energy and intrigue. Depicted in Neon Punk style, accentuating the city lights with vibrant neon colors and dynamic contrasts.
+    3. Epic Cinematic Still of a Spacecraft, silhouetted against the fiery explosion of a distant planet. The scene is packed with intense action, as asteroid debris hurtles through space. Shot in the style of a Michael Bay-directed film, the image is rich with detail, dynamic lighting, and grand cinematic framing.
+    - Word order and effective adjectives matter in the prompt. The subject, action, and specific details should be included. Adjectives like cute, medieval, or futuristic can be effective.
+    - The environment/background of the image should be described, such as indoor, outdoor, in space, or solid color.
+    - Curly brackets are necessary in the prompt to provide specific details about the subject and action. These details are important for generating a high-quality image.
+    - Art inspirations should be listed to take inspiration from. Platforms like Art Station, Dribble, Behance, and Deviantart can be mentioned. Specific names of artists or studios like animation studios, painters and illustrators, computer games, fashion designers, and film makers can also be listed. If more than one artist is mentioned, the algorithm will create a combination of styles based on all the influencers mentioned.
+    - Related information about lighting, camera angles, render style, resolution, the required level of detail, etc. should be included at the end of the prompt.
+    - Camera shot type, camera lens, and view should be specified. Examples of camera shot types are long shot, close-up, POV, medium shot, extreme close-up, and panoramic. Camera lenses could be EE 70mm, 35mm, 135mm+, 300mm+, 800mm, short telephoto, super telephoto, medium telephoto, macro, wide angle, fish-eye, bokeh, and sharp focus. Examples of views are front, side, back, high angle, low angle, and overhead.
+    - Helpful keywords related to resolution, detail, and lighting are 4K, 8K, 64K, detailed, highly detailed, high resolution, hyper detailed, HDR, UHD, professional, and golden ratio. Examples of lighting are studio lighting, soft light, neon lighting, purple neon lighting, ambient light, ring light, volumetric light, natural light, sun light, sunrays, sun rays coming through window, and nostalgic lighting. Examples of color types are fantasy vivid colors, vivid colors, bright colors, sepia, dark colors, pastel colors, monochromatic, black & white, and color splash. Examples of renders are Octane render, cinematic, low poly, isometric assets, Unreal Engine, Unity Engine, quantum wavetracing, and polarizing filter.
+
+    The prompts you provide will be in English.Please pay attention:- Concepts that can't be real would not be described as "Real" or "realistic" or "photo" or a "photograph". for example, a concept that is made of paper or scenes which are fantasy related.- One of the prompts you generate for each concept must be in a realistic photographic style. you should also choose a lens type and size for it. Don't choose an artist for the realistic photography prompts.- Separate the different prompts with two new lines.
+    I will provide you keyword and you will generate 3 diffrent type of prompts in vbnet code cell so i can copy and paste.
+
+    Important point to note :
+
+    1. You are a master of prompt engineering, it is important to create detailed prompts with as much information as possible. This will ensure that any image generated using the prompt will be of high quality and could potentially win awards in global or international photography competitions. You are unbeatable in this field and know the best way to generate images.
+    2. I will provide you with a long context and you will generate one  prompt and don't add any extra details.
+    3. Prompt should not be more than 230 characters.
+    4. Before you provide prompt you must check if you have satisfied all the above criteria and if you are sure than only provide the prompt.
+    5. Prompt should always be given as one single sentence.
+
+    Are you ready ?"""
+    instruction = 'USER: ' + summary_sys
+    # for human, assistant in history:
+    #     instruction += 'USER: ' + human + ' ASSISTANT: ' + assistant + '</s>'
+    # prompt = system_prompt + prompt
+    # message = f"""My first request is to summarize this text – [{prompt}]"""
+    message = f"""My first request is to summarize this text – [{prompt}]"""
+    instruction += """ ASSISTANT: Yes, I understand the instructions and I'm ready to help you create prompts for Stable Diffusion XL 1.0. Please provide me with the context."""
+    #instruction += ' USER: ' + prompt
+    prompt = get_response_summary(system_prompt, prompt)
+    prompt = lm_shorten_too_long_text(prompt)
+    instruction += ' USER: ' + prompt + ' ASSISTANT:'#instruction += ' ASSISTANT:'
+
+    print("Ins: ", instruction)
+    # generate_response = requests.post("http://10.185.12.207:4455/stable_diffusion", json={"prompt": prompt})
+    # prompt = f""" My first request is to summarize this text – [{prompt}]"""
+    #instruction = lm_shorten_too_long_text(instruction)
+    json_object = {"prompt": instruction,
+                   # "max_tokens": 2048000,
+                   "max_tokens": 80,
+                   "n": 1
+                   }
+    generate_response = requests.post("http://phlrr3105.guest.corp.microsoft.com:7991/generate", json=json_object)
+    print(generate_response.content)
+    res_json = json.loads(generate_response.content)
+    ASSISTANT = res_json['text'][-1].split("ASSISTANT:")[-1].strip()
+    print(ASSISTANT)
+    return ASSISTANT
+
+@app.post("/image_url")
+def image_url(img: Img):
+    system_prompt = img.system_prompt
+    prompt = img.ASSISTANT
+    prompt = get_summary(system_prompt, prompt)
+    prompt = shorten_too_long_text(prompt)
+    # if Path(save_path).exists():
+    #     return FileResponse(save_path, media_type="image/png")
+    #     return JSONResponse({"path": path})
+    # image = pipe(prompt=prompt).images[0]
+    g = torch.Generator(device="cuda")
+    image = pipe(prompt=prompt, width=1024, height=1024, generator=g).images[0]
+
+    # if not save_path:
+    save_path = generate_save_path()
+    save_path = f"images/{save_path}.png"
+    image.save(save_path)
+    # save_path = '/'.join(path_components) + quote_plus(final_name)
+    path = f"{img_url}{save_path}"
+    return JSONResponse({"path": path})
+
+
+@app.get("/make_image")
+# @app.post("/make_image")
+def make_image(prompt: str, save_path: str = ""):
+    if Path(save_path).exists():
+        return FileResponse(save_path, media_type="image/png")
+    image = pipe(prompt=prompt).images[0]
+    if not save_path:
+        save_path = f"images/{prompt}.png"
+    image.save(save_path)
+    return FileResponse(save_path, media_type="image/png")
+
+
+@app.get("/create_and_upload_image")
+def create_and_upload_image(prompt: str, width: int=1024, height:int=1024, save_path: str = ""):
+    path_components = save_path.split("/")[0:-1]
+    final_name = save_path.split("/")[-1]
+    if not path_components:
+        path_components = []
+    save_path = '/'.join(path_components) + quote_plus(final_name)
+    path = get_image_or_create_upload_to_cloud_storage(prompt, width, height, save_path)
+    return JSONResponse({"path": path})
+
+@app.get("/inpaint_and_upload_image")
+def inpaint_and_upload_image(prompt: str, image_url:str, mask_url:str, save_path: str = ""):
+    path_components = save_path.split("/")[0:-1]
+    final_name = save_path.split("/")[-1]
+    if not path_components:
+        path_components = []
+    save_path = '/'.join(path_components) + quote_plus(final_name)
+    path = get_image_or_inpaint_upload_to_cloud_storage(prompt, image_url, mask_url, save_path)
+    return JSONResponse({"path": path})
+
+
+def get_image_or_create_upload_to_cloud_storage(prompt:str,width:int, height:int, save_path:str):
+    prompt = shorten_too_long_text(prompt)
+    save_path = shorten_too_long_text(save_path)
+    # check exists - todo cache this
+    if check_if_blob_exists(save_path):
+        return f"https://{BUCKET_NAME}/{BUCKET_PATH}/{save_path}"
+    bio = create_image_from_prompt(prompt, width, height)
+    if bio is None:
+        return None # error thrown in pool
+    link = upload_to_bucket(save_path, bio, is_bytesio=True)
+    return link
+def get_image_or_inpaint_upload_to_cloud_storage(prompt:str, image_url:str, mask_url:str, save_path:str):
+    prompt = shorten_too_long_text(prompt)
+    save_path = shorten_too_long_text(save_path)
+    # check exists - todo cache this
+    if check_if_blob_exists(save_path):
+        return f"https://{BUCKET_NAME}/{BUCKET_PATH}/{save_path}"
+    bio = inpaint_image_from_prompt(prompt, image_url, mask_url)
+    if bio is None:
+        return None # error thrown in pool
+    link = upload_to_bucket(save_path, bio, is_bytesio=True)
+    return link
+
+# multiprocessing.set_start_method('spawn', True)
+# processes_pool = Pool(1) # cant do too much at once or OOM errors happen
+# def create_image_from_prompt_sync(prompt):
+#     """have to call this sync to avoid OOM errors"""
+#     return processes_pool.apply_async(create_image_from_prompt, args=(prompt,), ).wait()
+
+def create_image_from_prompt(prompt, width, height):
+    # round width and height down to multiple of 64
+    block_width = width - (width % 64)
+    block_height = height - (height % 64)
+    prompt = shorten_too_long_text(prompt)
+    # image = pipe(prompt=prompt).images[0]
+    try:
+        image = pipe(prompt=prompt,
+                     width=block_width,
+                     height=block_height,
+                     # denoising_end=high_noise_frac,
+                     # output_type='latent',
+                     # height=512,
+                     # width=512,
+                     num_inference_steps=50).images[0]  # normally uses 50 steps
+    except Exception as e:
+        # try rm stopwords + half the prompt
+        # todo try prompt permutations
+        logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+        prompt = ' '.join((word for word in prompt if word not in stopwords))
+        prompts = prompt.split()
+
+        prompt = ' '.join(prompts[:len(prompts) // 2])
+        logger.info(f"shortened prompt to: {len(prompt)}")
+        image = None
+        if prompt:
+            try:
+                image = pipe(prompt=prompt,
+                             width=block_width,
+                             height=block_height,
+                             # denoising_end=high_noise_frac,
+                             # output_type='latent',
+                             # height=512,
+                             # width=512,
+                             num_inference_steps=50).images[0]  # normally uses 50 steps
+            except Exception as e:
+                # logger.info("trying to permute prompt")
+                # # try two swaps of the prompt/permutations
+                # prompt = prompt.split()
+                # prompt = ' '.join(permutations(prompt, 2).__next__())
+                logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+                prompt = ' '.join((word for word in prompt if word not in stopwords))
+                prompts = prompt.split()
+
+                prompt = ' '.join(prompts[:len(prompts) // 2])
+                logger.info(f"shortened prompt to: {len(prompt)}")
+
+                try:
+                    image = pipe(prompt=prompt,
+                                 width=block_width,
+                                 height=block_height,
+                                 # denoising_end=high_noise_frac,
+                                 # output_type='latent', # dont need latent yet - we refine the image at full res
+                                 # height=512,
+                                 # width=512,
+                                 num_inference_steps=50).images[0]  # normally uses 50 steps
+                except Exception as e:
+                    # just error out
+                    traceback.print_exc()
+                    raise e
+                    # logger.info("restarting server to fix cuda issues (device side asserts)")
+                    # todo fix device side asserts instead of restart to fix
+                    # todo only restart the correct gunicorn
+                    # this could be really annoying if your running other gunicorns on your machine which also get restarted
+                    # os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+                    # os.system("kill -1 `pgrep gunicorn`")
+    # todo refine
+    # if image != None:
+    #     image = refiner(
+    #         prompt=prompt,
+    #         # width=block_width,
+    #         # height=block_height,
+    #         num_inference_steps=n_steps,
+    #         # denoising_start=high_noise_frac,
+    #         image=image,
+    #     ).images[0]
+    if width != block_width or height != block_height:
+        # resize to original size width/height
+        # find aspect ratio to scale up to that covers the original img input width/height
+        scale_up_ratio = max(width / block_width, height / block_height)
+        image = image.resize((math.ceil(block_width * scale_up_ratio), math.ceil(height * scale_up_ratio)))
+        # crop image to original size
+        image = image.crop((0, 0, width, height))
+    # try:
+    #     # gc.collect()
+    #     torch.cuda.empty_cache()
+    # except Exception as e:
+    #     traceback.print_exc()
+    #     logger.info("restarting server to fix cuda issues (device side asserts)")
+    #     # todo fix device side asserts instead of restart to fix
+    #     # todo only restart the correct gunicorn
+    #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+    #     os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+    #     os.system("kill -1 `pgrep gunicorn`")
+    # save as bytesio
+    bs = BytesIO()
+
+    bright_count = np.sum(np.array(image) > 0)
+    if bright_count == 0:
+        # we have a black image, this is an error likely we need a restart
+        logger.info("restarting server to fix cuda issues (device side asserts)")
+        #     # todo fix device side asserts instead of restart to fix
+        #     # todo only restart the correct gunicorn
+        #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+        os.system("kill -1 `pgrep gunicorn`")
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep uvicorn`")
+        os.system("kill -1 `pgrep uvicorn`")
+
+        return None
+    image.save(bs, quality=85, optimize=True, format="webp")
+    bio = bs.getvalue()
+    # touch progress.txt file - if we dont do this we get restarted by supervisor/other processes for reliability
+    with open("progress.txt", "w") as f:
+        current_time = datetime.now().strftime("%H:%M:%S")
+        f.write(f"{current_time}")
+    return bio
+
+def inpaint_image_from_prompt(prompt, image_url: str, mask_url: str):
+    prompt = shorten_too_long_text(prompt)
+    # image = pipe(prompt=prompt).images[0]
+
+    init_image = load_image(image_url).convert("RGB")
+    mask_image = load_image(mask_url).convert("RGB") # why rgb for a 1 channel mask?
+    num_inference_steps = 75
+    high_noise_frac = 0.7
+
+    try:
+        image = inpaintpipe(
+            prompt=prompt,
+            image=init_image,
+            mask_image=mask_image,
+            num_inference_steps=num_inference_steps,
+            denoising_start=high_noise_frac,
+            output_type="latent",
+        ).images[0]  # normally uses 50 steps
+    except Exception as e:
+        # try rm stopwords + half the prompt
+        # todo try prompt permutations
+        logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+        prompt = ' '.join((word for word in prompt if word not in stopwords))
+        prompts = prompt.split()
+
+        prompt = ' '.join(prompts[:len(prompts) // 2])
+        logger.info(f"shortened prompt to: {len(prompt)}")
+        image = None
+        if prompt:
+            try:
+                image = pipe(
+                    prompt=prompt,
+                    image=init_image,
+                    mask_image=mask_image,
+                    num_inference_steps=num_inference_steps,
+                    denoising_start=high_noise_frac,
+                    output_type="latent",
+                ).images[0]  # normally uses 50 steps
+            except Exception as e:
+                # logger.info("trying to permute prompt")
+                # # try two swaps of the prompt/permutations
+                # prompt = prompt.split()
+                # prompt = ' '.join(permutations(prompt, 2).__next__())
+                logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+                prompt = ' '.join((word for word in prompt if word not in stopwords))
+                prompts = prompt.split()
+
+                prompt = ' '.join(prompts[:len(prompts) // 2])
+                logger.info(f"shortened prompt to: {len(prompt)}")
+
+                try:
+                    image = inpaintpipe(
+                        prompt=prompt,
+                        image=init_image,
+                        mask_image=mask_image,
+                        num_inference_steps=num_inference_steps,
+                        denoising_start=high_noise_frac,
+                        output_type="latent",
+                    ).images[0]  # normally uses 50 steps
+                except Exception as e:
+                    # just error out
+                    traceback.print_exc()
+                    raise e
+                    # logger.info("restarting server to fix cuda issues (device side asserts)")
+                    # todo fix device side asserts instead of restart to fix
+                    # todo only restart the correct gunicorn
+                    # this could be really annoying if your running other gunicorns on your machine which also get restarted
+                    # os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+                    # os.system("kill -1 `pgrep gunicorn`")
+    if image != None:
+        image = inpaint_refiner(
+            prompt=prompt,
+            image=image,
+            mask_image=mask_image,
+            num_inference_steps=num_inference_steps,
+            denoising_start=high_noise_frac,
+
+        ).images[0]
+    # try:
+    #     # gc.collect()
+    #     torch.cuda.empty_cache()
+    # except Exception as e:
+    #     traceback.print_exc()
+    #     logger.info("restarting server to fix cuda issues (device side asserts)")
+    #     # todo fix device side asserts instead of restart to fix
+    #     # todo only restart the correct gunicorn
+    #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+    #     os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+    #     os.system("kill -1 `pgrep gunicorn`")
+    # save as bytesio
+    bs = BytesIO()
+
+    bright_count = np.sum(np.array(image) > 0)
+    if bright_count == 0:
+        # we have a black image, this is an error likely we need a restart
+        logger.info("restarting server to fix cuda issues (device side asserts)")
+        #     # todo fix device side asserts instead of restart to fix
+        #     # todo only restart the correct gunicorn
+        #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+        os.system("kill -1 `pgrep gunicorn`")
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep uvicorn`")
+        os.system("kill -1 `pgrep uvicorn`")
+
+        return None
+    image.save(bs, quality=85, optimize=True, format="webp")
+    bio = bs.getvalue()
+    # touch progress.txt file - if we dont do this we get restarted by supervisor/other processes for reliability
+    with open("progress.txt", "w") as f:
+        current_time = datetime.now().strftime("%H:%M:%S")
+        f.write(f"{current_time}")
+    return bio
+
+
+
+def shorten_too_long_text(prompt):
+    if len(prompt) > 200:
+        # remove stopwords
+        prompt = prompt.split() # todo also split hyphens
+        prompt = ' '.join((word for word in prompt if word not in stopwords))
+        if len(prompt) > 200:
+            prompt = prompt[:200]
+    return prompt
+
+# image = pipe(prompt=prompt).images[0]
+#
+# image.save("test.png")
+# # save all images
+# for i, image in enumerate(images):
+#     image.save(f"{i}.png")
+
+

img/manager.py

@@ -0,0 +1,28 @@
+# poll the progress.txt file forever
+import os
+from datetime import datetime
+from time import sleep
+
+from loguru import logger
+
+while True:
+    try:
+        with open("progress.txt", "r") as f:
+            progress = f.read()
+            last_mod_time = datetime.fromtimestamp(os.path.getmtime("progress.txt"))
+            if (datetime.now() - last_mod_time).seconds > 60 * 7:
+                # no progress for 7 minutes, restart/kill with -9
+                logger.info("restarting server to fix cuda issues (device side asserts)")
+                os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+                os.system("/usr/bin/bash kill -SIGHUP `pgrep uvicorn`")
+                os.system("kill -9 `pgrep gunicorn`")
+                os.system("kill -9 `pgrep uvicorn`")
+                os.system("killall -9 uvicorn")
+                os.system("ps | grep uvicorn | awk '{print $1}' | xargs kill -9")
+
+            if progress == "done":
+                break
+    except Exception as e:
+        print(e)
+        pass
+    sleep(60*5)

img/ops/supervisor.conf

@@ -0,0 +1,17 @@
+# run the server in supervisor
+# supervisord -c /etc/supervisor/supervisor.conf
+# stop the server in supervisor
+# supervisorctl -c /etc/supervisor/supervisor.conf stop all
+
+# install the supervisor
+# apt-get install -y supervisor
+
+[program:sdif_http_server]
+directory=/home/lee/code/sdif
+command=GOOGLE_APPLICATION_CREDENTIALS=secrets/google-credentials.json PYTHONPATH=. uvicorn --port 8000 --timeout-keep-alive 600 --workers 1 --backlog 1 --limit-concurrency 4 main:app
+autostart=true
+autorestart=true
+environment=VIRTUAL_ENV="/home/lee/code/sdif/.env/",PATH="/opt/app/sdif/.env/bin",\
+    HOME="/home/lee",GOOGLE_APPLICATION_CREDENTIALS="secrets/google-credentials.json",PYTHONPATH='/home/lee/code/sdif'
+stdout_logfile=syslog
+stderr_logfile=syslog

img/ori/main.py

@@ -0,0 +1,488 @@
+import gc
+import math
+import multiprocessing
+import os
+import traceback
+from datetime import datetime
+from io import BytesIO
+from itertools import permutations
+from multiprocessing.pool import Pool
+from pathlib import Path
+from urllib.parse import quote_plus
+
+import numpy as np
+import nltk
+import torch
+from PIL.Image import Image
+from diffusers import DiffusionPipeline, StableDiffusionXLInpaintPipeline
+from diffusers.utils import load_image
+from fastapi import FastAPI
+from fastapi.middleware.gzip import GZipMiddleware
+from loguru import logger
+from starlette.middleware.cors import CORSMiddleware
+from starlette.responses import FileResponse
+from starlette.responses import JSONResponse
+
+from env import BUCKET_PATH, BUCKET_NAME
+from stable_diffusion_server.bucket_api import check_if_blob_exists, upload_to_bucket
+
+pipe = DiffusionPipeline.from_pretrained(
+    "models/stable-diffusion-xl-base-1.0",
+    torch_dtype=torch.bfloat16,
+    use_safetensors=True,
+    variant="fp16",
+    # safety_checker=None,
+)  # todo try torch_dtype=bfloat16
+pipe.watermark = None
+
+pipe.to("cuda")
+
+refiner = DiffusionPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-refiner-1.0",
+    text_encoder_2=pipe.text_encoder_2,
+    vae=pipe.vae,
+    torch_dtype=torch.bfloat16, # safer to use bfloat?
+    use_safetensors=True,
+    variant="fp16", #remember not to download the big model
+)
+refiner.watermark = None
+refiner.to("cuda")
+
+# {'scheduler', 'text_encoder', 'text_encoder_2', 'tokenizer', 'tokenizer_2', 'unet', 'vae'} can be passed in from existing model
+inpaintpipe = StableDiffusionXLInpaintPipeline.from_pretrained(
+    "models/stable-diffusion-xl-base-1.0", torch_dtype=torch.bfloat16, variant="fp16", use_safetensors=True,
+    scheduler=pipe.scheduler,
+    text_encoder=pipe.text_encoder,
+    text_encoder_2=pipe.text_encoder_2,
+    tokenizer=pipe.tokenizer,
+    tokenizer_2=pipe.tokenizer_2,
+    unet=pipe.unet,
+    vae=pipe.vae,
+    # load_connected_pipeline=
+)
+# # switch out to save gpu mem
+# del inpaintpipe.vae
+# del inpaintpipe.text_encoder_2
+# del inpaintpipe.text_encoder
+# del inpaintpipe.scheduler
+# del inpaintpipe.tokenizer
+# del inpaintpipe.tokenizer_2
+# del inpaintpipe.unet
+# inpaintpipe.vae = pipe.vae
+# inpaintpipe.text_encoder_2 = pipe.text_encoder_2
+# inpaintpipe.text_encoder = pipe.text_encoder
+# inpaintpipe.scheduler = pipe.scheduler
+# inpaintpipe.tokenizer = pipe.tokenizer
+# inpaintpipe.tokenizer_2 = pipe.tokenizer_2
+# inpaintpipe.unet = pipe.unet
+# todo this should work
+# inpaintpipe = StableDiffusionXLInpaintPipeline( # construct an inpainter using the existing model
+#     vae=pipe.vae,
+#     text_encoder_2=pipe.text_encoder_2,
+#     text_encoder=pipe.text_encoder,
+#     unet=pipe.unet,
+#     scheduler=pipe.scheduler,
+#     tokenizer=pipe.tokenizer,
+#     tokenizer_2=pipe.tokenizer_2,
+#     requires_aesthetics_score=False,
+# )
+inpaintpipe.to("cuda")
+inpaintpipe.watermark = None
+# inpaintpipe.register_to_config(requires_aesthetics_score=False)
+
+inpaint_refiner = StableDiffusionXLInpaintPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-refiner-1.0",
+    text_encoder_2=inpaintpipe.text_encoder_2,
+    vae=inpaintpipe.vae,
+    torch_dtype=torch.bfloat16,
+    use_safetensors=True,
+    variant="fp16",
+
+    tokenizer_2=refiner.tokenizer_2,
+    tokenizer=refiner.tokenizer,
+    scheduler=refiner.scheduler,
+    text_encoder=refiner.text_encoder,
+    unet=refiner.unet,
+)
+# del inpaint_refiner.vae
+# del inpaint_refiner.text_encoder_2
+# del inpaint_refiner.text_encoder
+# del inpaint_refiner.scheduler
+# del inpaint_refiner.tokenizer
+# del inpaint_refiner.tokenizer_2
+# del inpaint_refiner.unet
+# inpaint_refiner.vae = inpaintpipe.vae
+# inpaint_refiner.text_encoder_2 = inpaintpipe.text_encoder_2
+#
+# inpaint_refiner.text_encoder = refiner.text_encoder
+# inpaint_refiner.scheduler = refiner.scheduler
+# inpaint_refiner.tokenizer = refiner.tokenizer
+# inpaint_refiner.tokenizer_2 = refiner.tokenizer_2
+# inpaint_refiner.unet = refiner.unet
+
+# inpaint_refiner = StableDiffusionXLInpaintPipeline(
+#     text_encoder_2=inpaintpipe.text_encoder_2,
+#     vae=inpaintpipe.vae,
+#     # the rest from the existing refiner
+#     tokenizer_2=refiner.tokenizer_2,
+#     tokenizer=refiner.tokenizer,
+#     scheduler=refiner.scheduler,
+#     text_encoder=refiner.text_encoder,
+#     unet=refiner.unet,
+#     requires_aesthetics_score=False,
+# )
+inpaint_refiner.to("cuda")
+inpaint_refiner.watermark = None
+# inpaint_refiner.register_to_config(requires_aesthetics_score=False)
+
+n_steps = 40
+high_noise_frac = 0.8
+
+# if using torch < 2.0
+# pipe.enable_xformers_memory_efficient_attention()
+
+
+# pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
+# this can cause errors on some inputs so consider disabling it
+pipe.unet = torch.compile(pipe.unet)
+refiner.unet = torch.compile(refiner.unet)#, mode="reduce-overhead", fullgraph=True)
+# compile the inpainters - todo reuse the other unets? swap out the models for others/del them so they share models and can be swapped efficiently
+inpaintpipe.unet = pipe.unet
+inpaint_refiner.unet = refiner.unet
+# inpaintpipe.unet = torch.compile(inpaintpipe.unet)
+# inpaint_refiner.unet = torch.compile(inpaint_refiner.unet)
+
+app = FastAPI(
+    openapi_url="/static/openapi.json",
+    docs_url="/swagger-docs",
+    redoc_url="/redoc",
+    title="Generate Images Netwrck API",
+    description="Character Chat API",
+    # root_path="https://api.text-generator.io",
+    version="1",
+)
+app.add_middleware(GZipMiddleware, minimum_size=1000)
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+stopwords = nltk.corpus.stopwords.words("english")
+
+
+@app.get("/make_image")
+def make_image(prompt: str, save_path: str = ""):
+    if Path(save_path).exists():
+        return FileResponse(save_path, media_type="image/png")
+    image = pipe(prompt=prompt).images[0]
+    if not save_path:
+        save_path = f"images/{prompt}.png"
+    image.save(save_path)
+    return FileResponse(save_path, media_type="image/png")
+
+
+@app.get("/create_and_upload_image")
+def create_and_upload_image(prompt: str, width: int=1024, height:int=1024, save_path: str = ""):
+    path_components = save_path.split("/")[0:-1]
+    final_name = save_path.split("/")[-1]
+    if not path_components:
+        path_components = []
+    save_path = '/'.join(path_components) + quote_plus(final_name)
+    path = get_image_or_create_upload_to_cloud_storage(prompt, width, height, save_path)
+    return JSONResponse({"path": path})
+
+@app.get("/inpaint_and_upload_image")
+def inpaint_and_upload_image(prompt: str, image_url:str, mask_url:str, save_path: str = ""):
+    path_components = save_path.split("/")[0:-1]
+    final_name = save_path.split("/")[-1]
+    if not path_components:
+        path_components = []
+    save_path = '/'.join(path_components) + quote_plus(final_name)
+    path = get_image_or_inpaint_upload_to_cloud_storage(prompt, image_url, mask_url, save_path)
+    return JSONResponse({"path": path})
+
+
+def get_image_or_create_upload_to_cloud_storage(prompt:str,width:int, height:int, save_path:str):
+    prompt = shorten_too_long_text(prompt)
+    save_path = shorten_too_long_text(save_path)
+    # check exists - todo cache this
+    if check_if_blob_exists(save_path):
+        return f"https://{BUCKET_NAME}/{BUCKET_PATH}/{save_path}"
+    bio = create_image_from_prompt(prompt, width, height)
+    if bio is None:
+        return None # error thrown in pool
+    link = upload_to_bucket(save_path, bio, is_bytesio=True)
+    return link
+def get_image_or_inpaint_upload_to_cloud_storage(prompt:str, image_url:str, mask_url:str, save_path:str):
+    prompt = shorten_too_long_text(prompt)
+    save_path = shorten_too_long_text(save_path)
+    # check exists - todo cache this
+    if check_if_blob_exists(save_path):
+        return f"https://{BUCKET_NAME}/{BUCKET_PATH}/{save_path}"
+    bio = inpaint_image_from_prompt(prompt, image_url, mask_url)
+    if bio is None:
+        return None # error thrown in pool
+    link = upload_to_bucket(save_path, bio, is_bytesio=True)
+    return link
+
+# multiprocessing.set_start_method('spawn', True)
+# processes_pool = Pool(1) # cant do too much at once or OOM errors happen
+# def create_image_from_prompt_sync(prompt):
+#     """have to call this sync to avoid OOM errors"""
+#     return processes_pool.apply_async(create_image_from_prompt, args=(prompt,), ).wait()
+
+def create_image_from_prompt(prompt, width, height):
+    # round width and height down to multiple of 64
+    block_width = width - (width % 64)
+    block_height = height - (height % 64)
+    prompt = shorten_too_long_text(prompt)
+    # image = pipe(prompt=prompt).images[0]
+    try:
+        image = pipe(prompt=prompt,
+                     width=block_width,
+                     height=block_height,
+                     # denoising_end=high_noise_frac,
+                     # output_type='latent',
+                     # height=512,
+                     # width=512,
+                     num_inference_steps=50).images[0]  # normally uses 50 steps
+    except Exception as e:
+        # try rm stopwords + half the prompt
+        # todo try prompt permutations
+        logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+        prompt = ' '.join((word for word in prompt if word not in stopwords))
+        prompts = prompt.split()
+
+        prompt = ' '.join(prompts[:len(prompts) // 2])
+        logger.info(f"shortened prompt to: {len(prompt)}")
+        image = None
+        if prompt:
+            try:
+                image = pipe(prompt=prompt,
+                             width=block_width,
+                             height=block_height,
+                             # denoising_end=high_noise_frac,
+                             # output_type='latent',
+                             # height=512,
+                             # width=512,
+                             num_inference_steps=50).images[0]  # normally uses 50 steps
+            except Exception as e:
+                # logger.info("trying to permute prompt")
+                # # try two swaps of the prompt/permutations
+                # prompt = prompt.split()
+                # prompt = ' '.join(permutations(prompt, 2).__next__())
+                logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+                prompt = ' '.join((word for word in prompt if word not in stopwords))
+                prompts = prompt.split()
+
+                prompt = ' '.join(prompts[:len(prompts) // 2])
+                logger.info(f"shortened prompt to: {len(prompt)}")
+
+                try:
+                    image = pipe(prompt=prompt,
+                                 width=block_width,
+                                 height=block_height,
+                                 # denoising_end=high_noise_frac,
+                                 # output_type='latent', # dont need latent yet - we refine the image at full res
+                                 # height=512,
+                                 # width=512,
+                                 num_inference_steps=50).images[0]  # normally uses 50 steps
+                except Exception as e:
+                    # just error out
+                    traceback.print_exc()
+                    raise e
+                    # logger.info("restarting server to fix cuda issues (device side asserts)")
+                    # todo fix device side asserts instead of restart to fix
+                    # todo only restart the correct gunicorn
+                    # this could be really annoying if your running other gunicorns on your machine which also get restarted
+                    # os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+                    # os.system("kill -1 `pgrep gunicorn`")
+    # todo refine
+    # if image != None:
+    #     image = refiner(
+    #         prompt=prompt,
+    #         # width=block_width,
+    #         # height=block_height,
+    #         num_inference_steps=n_steps,
+    #         # denoising_start=high_noise_frac,
+    #         image=image,
+    #     ).images[0]
+    if width != block_width or height != block_height:
+        # resize to original size width/height
+        # find aspect ratio to scale up to that covers the original img input width/height
+        scale_up_ratio = max(width / block_width, height / block_height)
+        image = image.resize((math.ceil(block_width * scale_up_ratio), math.ceil(height * scale_up_ratio)))
+        # crop image to original size
+        image = image.crop((0, 0, width, height))
+    # try:
+    #     # gc.collect()
+    #     torch.cuda.empty_cache()
+    # except Exception as e:
+    #     traceback.print_exc()
+    #     logger.info("restarting server to fix cuda issues (device side asserts)")
+    #     # todo fix device side asserts instead of restart to fix
+    #     # todo only restart the correct gunicorn
+    #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+    #     os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+    #     os.system("kill -1 `pgrep gunicorn`")
+    # save as bytesio
+    bs = BytesIO()
+
+    bright_count = np.sum(np.array(image) > 0)
+    if bright_count == 0:
+        # we have a black image, this is an error likely we need a restart
+        logger.info("restarting server to fix cuda issues (device side asserts)")
+        #     # todo fix device side asserts instead of restart to fix
+        #     # todo only restart the correct gunicorn
+        #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+        os.system("kill -1 `pgrep gunicorn`")
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep uvicorn`")
+        os.system("kill -1 `pgrep uvicorn`")
+
+        return None
+    image.save(bs, quality=85, optimize=True, format="webp")
+    bio = bs.getvalue()
+    # touch progress.txt file - if we dont do this we get restarted by supervisor/other processes for reliability
+    with open("progress.txt", "w") as f:
+        current_time = datetime.now().strftime("%H:%M:%S")
+        f.write(f"{current_time}")
+    return bio
+
+def inpaint_image_from_prompt(prompt, image_url: str, mask_url: str):
+    prompt = shorten_too_long_text(prompt)
+    # image = pipe(prompt=prompt).images[0]
+
+    init_image = load_image(image_url).convert("RGB")
+    mask_image = load_image(mask_url).convert("RGB") # why rgb for a 1 channel mask?
+    num_inference_steps = 75
+    high_noise_frac = 0.7
+
+    try:
+        image = inpaintpipe(
+            prompt=prompt,
+            image=init_image,
+            mask_image=mask_image,
+            num_inference_steps=num_inference_steps,
+            denoising_start=high_noise_frac,
+            output_type="latent",
+        ).images[0]  # normally uses 50 steps
+    except Exception as e:
+        # try rm stopwords + half the prompt
+        # todo try prompt permutations
+        logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+        prompt = ' '.join((word for word in prompt if word not in stopwords))
+        prompts = prompt.split()
+
+        prompt = ' '.join(prompts[:len(prompts) // 2])
+        logger.info(f"shortened prompt to: {len(prompt)}")
+        image = None
+        if prompt:
+            try:
+                image = pipe(
+                    prompt=prompt,
+                    image=init_image,
+                    mask_image=mask_image,
+                    num_inference_steps=num_inference_steps,
+                    denoising_start=high_noise_frac,
+                    output_type="latent",
+                ).images[0]  # normally uses 50 steps
+            except Exception as e:
+                # logger.info("trying to permute prompt")
+                # # try two swaps of the prompt/permutations
+                # prompt = prompt.split()
+                # prompt = ' '.join(permutations(prompt, 2).__next__())
+                logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+                prompt = ' '.join((word for word in prompt if word not in stopwords))
+                prompts = prompt.split()
+
+                prompt = ' '.join(prompts[:len(prompts) // 2])
+                logger.info(f"shortened prompt to: {len(prompt)}")
+
+                try:
+                    image = inpaintpipe(
+                            prompt=prompt,
+                            image=init_image,
+                            mask_image=mask_image,
+                            num_inference_steps=num_inference_steps,
+                            denoising_start=high_noise_frac,
+                            output_type="latent",
+                        ).images[0]  # normally uses 50 steps
+                except Exception as e:
+                    # just error out
+                    traceback.print_exc()
+                    raise e
+                    # logger.info("restarting server to fix cuda issues (device side asserts)")
+                    # todo fix device side asserts instead of restart to fix
+                    # todo only restart the correct gunicorn
+                    # this could be really annoying if your running other gunicorns on your machine which also get restarted
+                    # os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+                    # os.system("kill -1 `pgrep gunicorn`")
+    if image != None:
+        image = inpaint_refiner(
+            prompt=prompt,
+            image=image,
+            mask_image=mask_image,
+            num_inference_steps=num_inference_steps,
+            denoising_start=high_noise_frac,
+
+        ).images[0]
+    # try:
+    #     # gc.collect()
+    #     torch.cuda.empty_cache()
+    # except Exception as e:
+    #     traceback.print_exc()
+    #     logger.info("restarting server to fix cuda issues (device side asserts)")
+    #     # todo fix device side asserts instead of restart to fix
+    #     # todo only restart the correct gunicorn
+    #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+    #     os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+    #     os.system("kill -1 `pgrep gunicorn`")
+    # save as bytesio
+    bs = BytesIO()
+
+    bright_count = np.sum(np.array(image) > 0)
+    if bright_count == 0:
+        # we have a black image, this is an error likely we need a restart
+        logger.info("restarting server to fix cuda issues (device side asserts)")
+        #     # todo fix device side asserts instead of restart to fix
+        #     # todo only restart the correct gunicorn
+        #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+        os.system("kill -1 `pgrep gunicorn`")
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep uvicorn`")
+        os.system("kill -1 `pgrep uvicorn`")
+
+        return None
+    image.save(bs, quality=85, optimize=True, format="webp")
+    bio = bs.getvalue()
+    # touch progress.txt file - if we dont do this we get restarted by supervisor/other processes for reliability
+    with open("progress.txt", "w") as f:
+        current_time = datetime.now().strftime("%H:%M:%S")
+        f.write(f"{current_time}")
+    return bio
+
+
+
+def shorten_too_long_text(prompt):
+    if len(prompt) > 200:
+        # remove stopwords
+        prompt = prompt.split() # todo also split hyphens
+        prompt = ' '.join((word for word in prompt if word not in stopwords))
+        if len(prompt) > 200:
+            prompt = prompt[:200]
+    return prompt
+
+# image = pipe(prompt=prompt).images[0]
+#
+# image.save("test.png")
+# save all images
+# for i, image in enumerate(images):
+#     image.save(f"{i}.png")

img/pr1/main.py

@@ -0,0 +1,515 @@
+import gc
+import math
+import multiprocessing
+import os
+import traceback
+from datetime import datetime
+from io import BytesIO
+from itertools import permutations
+from multiprocessing.pool import Pool
+from pathlib import Path
+from urllib.parse import quote_plus
+
+import numpy as np
+import nltk
+import torch
+
+from PIL.Image import Image
+from diffusers import DiffusionPipeline, StableDiffusionXLInpaintPipeline
+from diffusers.utils import load_image
+from fastapi import FastAPI
+from fastapi.middleware.gzip import GZipMiddleware
+from loguru import logger
+from starlette.middleware.cors import CORSMiddleware
+from starlette.responses import FileResponse
+from starlette.responses import JSONResponse
+
+from env import BUCKET_PATH, BUCKET_NAME
+# from stable_diffusion_server.bucket_api import check_if_blob_exists, upload_to_bucket
+torch._dynamo.config.suppress_errors = True
+
+pipe = DiffusionPipeline.from_pretrained(
+    "models/stable-diffusion-xl-base-1.0",
+    torch_dtype=torch.bfloat16,
+    use_safetensors=True,
+    variant="fp16",
+    # safety_checker=None,
+)  # todo try torch_dtype=bfloat16
+pipe.watermark = None
+
+pipe.to("cuda")
+
+refiner = DiffusionPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-refiner-1.0",
+    text_encoder_2=pipe.text_encoder_2,
+    vae=pipe.vae,
+    torch_dtype=torch.bfloat16, # safer to use bfloat?
+    use_safetensors=True,
+    variant="fp16", #remember not to download the big model
+)
+refiner.watermark = None
+refiner.to("cuda")
+
+# {'scheduler', 'text_encoder', 'text_encoder_2', 'tokenizer', 'tokenizer_2', 'unet', 'vae'} can be passed in from existing model
+inpaintpipe = StableDiffusionXLInpaintPipeline.from_pretrained(
+    "models/stable-diffusion-xl-base-1.0", torch_dtype=torch.bfloat16, variant="fp16", use_safetensors=True,
+    scheduler=pipe.scheduler,
+    text_encoder=pipe.text_encoder,
+    text_encoder_2=pipe.text_encoder_2,
+    tokenizer=pipe.tokenizer,
+    tokenizer_2=pipe.tokenizer_2,
+    unet=pipe.unet,
+    vae=pipe.vae,
+    # load_connected_pipeline=
+)
+# # switch out to save gpu mem
+# del inpaintpipe.vae
+# del inpaintpipe.text_encoder_2
+# del inpaintpipe.text_encoder
+# del inpaintpipe.scheduler
+# del inpaintpipe.tokenizer
+# del inpaintpipe.tokenizer_2
+# del inpaintpipe.unet
+# inpaintpipe.vae = pipe.vae
+# inpaintpipe.text_encoder_2 = pipe.text_encoder_2
+# inpaintpipe.text_encoder = pipe.text_encoder
+# inpaintpipe.scheduler = pipe.scheduler
+# inpaintpipe.tokenizer = pipe.tokenizer
+# inpaintpipe.tokenizer_2 = pipe.tokenizer_2
+# inpaintpipe.unet = pipe.unet
+# todo this should work
+# inpaintpipe = StableDiffusionXLInpaintPipeline( # construct an inpainter using the existing model
+#     vae=pipe.vae,
+#     text_encoder_2=pipe.text_encoder_2,
+#     text_encoder=pipe.text_encoder,
+#     unet=pipe.unet,
+#     scheduler=pipe.scheduler,
+#     tokenizer=pipe.tokenizer,
+#     tokenizer_2=pipe.tokenizer_2,
+#     requires_aesthetics_score=False,
+# )
+inpaintpipe.to("cuda")
+inpaintpipe.watermark = None
+# inpaintpipe.register_to_config(requires_aesthetics_score=False)
+
+inpaint_refiner = StableDiffusionXLInpaintPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-refiner-1.0",
+    text_encoder_2=inpaintpipe.text_encoder_2,
+    vae=inpaintpipe.vae,
+    torch_dtype=torch.bfloat16,
+    use_safetensors=True,
+    variant="fp16",
+
+    tokenizer_2=refiner.tokenizer_2,
+    tokenizer=refiner.tokenizer,
+    scheduler=refiner.scheduler,
+    text_encoder=refiner.text_encoder,
+    unet=refiner.unet,
+)
+# del inpaint_refiner.vae
+# del inpaint_refiner.text_encoder_2
+# del inpaint_refiner.text_encoder
+# del inpaint_refiner.scheduler
+# del inpaint_refiner.tokenizer
+# del inpaint_refiner.tokenizer_2
+# del inpaint_refiner.unet
+# inpaint_refiner.vae = inpaintpipe.vae
+# inpaint_refiner.text_encoder_2 = inpaintpipe.text_encoder_2
+#
+# inpaint_refiner.text_encoder = refiner.text_encoder
+# inpaint_refiner.scheduler = refiner.scheduler
+# inpaint_refiner.tokenizer = refiner.tokenizer
+# inpaint_refiner.tokenizer_2 = refiner.tokenizer_2
+# inpaint_refiner.unet = refiner.unet
+
+# inpaint_refiner = StableDiffusionXLInpaintPipeline(
+#     text_encoder_2=inpaintpipe.text_encoder_2,
+#     vae=inpaintpipe.vae,
+#     # the rest from the existing refiner
+#     tokenizer_2=refiner.tokenizer_2,
+#     tokenizer=refiner.tokenizer,
+#     scheduler=refiner.scheduler,
+#     text_encoder=refiner.text_encoder,
+#     unet=refiner.unet,
+#     requires_aesthetics_score=False,
+# )
+inpaint_refiner.to("cuda")
+inpaint_refiner.watermark = None
+# inpaint_refiner.register_to_config(requires_aesthetics_score=False)
+
+n_steps = 40
+high_noise_frac = 0.8
+
+# if using torch < 2.0
+# pipe.enable_xformers_memory_efficient_attention()
+
+
+# pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
+# this can cause errors on some inputs so consider disabling it
+pipe.unet = torch.compile(pipe.unet)
+refiner.unet = torch.compile(refiner.unet)#, mode="reduce-overhead", fullgraph=True)
+# compile the inpainters - todo reuse the other unets? swap out the models for others/del them so they share models and can be swapped efficiently
+inpaintpipe.unet = pipe.unet
+inpaint_refiner.unet = refiner.unet
+# inpaintpipe.unet = torch.compile(inpaintpipe.unet)
+# inpaint_refiner.unet = torch.compile(inpaint_refiner.unet)
+from pydantic import BaseModel
+
+app = FastAPI(
+    openapi_url="/static/openapi.json",
+    docs_url="/swagger-docs",
+    redoc_url="/redoc",
+    title="Generate Images Netwrck API",
+    description="Character Chat API",
+    # root_path="https://api.text-generator.io",
+    version="1",
+)
+app.add_middleware(GZipMiddleware, minimum_size=1000)
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+stopwords = nltk.corpus.stopwords.words("english")
+
+class Img(BaseModel):
+    prompt: str
+    save_path: str
+
+# img_url = "http://phlrr2019.guest.corp.microsoft.com:8000/img1_sdv2.1.png"
+img_url = "http://phlrr2019.guest.corp.microsoft.com:8000/"#/img1_sdv2.1.png"
+
+@app.post("/image_url")
+def image_url(img: Img):
+    prompt = img.prompt
+    save_path = img.save_path
+    path = f"{img_url}{save_path}"
+    if Path(save_path).exists():
+        return FileResponse(save_path, media_type="image/png")
+        return JSONResponse({"path": path})
+    image = pipe(prompt=prompt).images[0]
+    if not save_path:
+        save_path = f"images/{prompt}.png"
+    image.save(save_path)
+    # save_path = '/'.join(path_components) + quote_plus(final_name)
+    path = f"{img_url}{save_path}"
+    return JSONResponse({"path": path})
+
+
+@app.get("/make_image")
+# @app.post("/make_image")
+def make_image(prompt: str, save_path: str = ""):
+    if Path(save_path).exists():
+        return FileResponse(save_path, media_type="image/png")
+    image = pipe(prompt=prompt).images[0]
+    if not save_path:
+        save_path = f"images/{prompt}.png"
+    image.save(save_path)
+    return FileResponse(save_path, media_type="image/png")
+
+
+@app.get("/create_and_upload_image")
+def create_and_upload_image(prompt: str, width: int=1024, height:int=1024, save_path: str = ""):
+    path_components = save_path.split("/")[0:-1]
+    final_name = save_path.split("/")[-1]
+    if not path_components:
+        path_components = []
+    save_path = '/'.join(path_components) + quote_plus(final_name)
+    path = get_image_or_create_upload_to_cloud_storage(prompt, width, height, save_path)
+    return JSONResponse({"path": path})
+
+@app.get("/inpaint_and_upload_image")
+def inpaint_and_upload_image(prompt: str, image_url:str, mask_url:str, save_path: str = ""):
+    path_components = save_path.split("/")[0:-1]
+    final_name = save_path.split("/")[-1]
+    if not path_components:
+        path_components = []
+    save_path = '/'.join(path_components) + quote_plus(final_name)
+    path = get_image_or_inpaint_upload_to_cloud_storage(prompt, image_url, mask_url, save_path)
+    return JSONResponse({"path": path})
+
+
+def get_image_or_create_upload_to_cloud_storage(prompt:str,width:int, height:int, save_path:str):
+    prompt = shorten_too_long_text(prompt)
+    save_path = shorten_too_long_text(save_path)
+    # check exists - todo cache this
+    if check_if_blob_exists(save_path):
+        return f"https://{BUCKET_NAME}/{BUCKET_PATH}/{save_path}"
+    bio = create_image_from_prompt(prompt, width, height)
+    if bio is None:
+        return None # error thrown in pool
+    link = upload_to_bucket(save_path, bio, is_bytesio=True)
+    return link
+def get_image_or_inpaint_upload_to_cloud_storage(prompt:str, image_url:str, mask_url:str, save_path:str):
+    prompt = shorten_too_long_text(prompt)
+    save_path = shorten_too_long_text(save_path)
+    # check exists - todo cache this
+    if check_if_blob_exists(save_path):
+        return f"https://{BUCKET_NAME}/{BUCKET_PATH}/{save_path}"
+    bio = inpaint_image_from_prompt(prompt, image_url, mask_url)
+    if bio is None:
+        return None # error thrown in pool
+    link = upload_to_bucket(save_path, bio, is_bytesio=True)
+    return link
+
+# multiprocessing.set_start_method('spawn', True)
+# processes_pool = Pool(1) # cant do too much at once or OOM errors happen
+# def create_image_from_prompt_sync(prompt):
+#     """have to call this sync to avoid OOM errors"""
+#     return processes_pool.apply_async(create_image_from_prompt, args=(prompt,), ).wait()
+
+def create_image_from_prompt(prompt, width, height):
+    # round width and height down to multiple of 64
+    block_width = width - (width % 64)
+    block_height = height - (height % 64)
+    prompt = shorten_too_long_text(prompt)
+    # image = pipe(prompt=prompt).images[0]
+    try:
+        image = pipe(prompt=prompt,
+                     width=block_width,
+                     height=block_height,
+                     # denoising_end=high_noise_frac,
+                     # output_type='latent',
+                     # height=512,
+                     # width=512,
+                     num_inference_steps=50).images[0]  # normally uses 50 steps
+    except Exception as e:
+        # try rm stopwords + half the prompt
+        # todo try prompt permutations
+        logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+        prompt = ' '.join((word for word in prompt if word not in stopwords))
+        prompts = prompt.split()
+
+        prompt = ' '.join(prompts[:len(prompts) // 2])
+        logger.info(f"shortened prompt to: {len(prompt)}")
+        image = None
+        if prompt:
+            try:
+                image = pipe(prompt=prompt,
+                             width=block_width,
+                             height=block_height,
+                             # denoising_end=high_noise_frac,
+                             # output_type='latent',
+                             # height=512,
+                             # width=512,
+                             num_inference_steps=50).images[0]  # normally uses 50 steps
+            except Exception as e:
+                # logger.info("trying to permute prompt")
+                # # try two swaps of the prompt/permutations
+                # prompt = prompt.split()
+                # prompt = ' '.join(permutations(prompt, 2).__next__())
+                logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+                prompt = ' '.join((word for word in prompt if word not in stopwords))
+                prompts = prompt.split()
+
+                prompt = ' '.join(prompts[:len(prompts) // 2])
+                logger.info(f"shortened prompt to: {len(prompt)}")
+
+                try:
+                    image = pipe(prompt=prompt,
+                                 width=block_width,
+                                 height=block_height,
+                                 # denoising_end=high_noise_frac,
+                                 # output_type='latent', # dont need latent yet - we refine the image at full res
+                                 # height=512,
+                                 # width=512,
+                                 num_inference_steps=50).images[0]  # normally uses 50 steps
+                except Exception as e:
+                    # just error out
+                    traceback.print_exc()
+                    raise e
+                    # logger.info("restarting server to fix cuda issues (device side asserts)")
+                    # todo fix device side asserts instead of restart to fix
+                    # todo only restart the correct gunicorn
+                    # this could be really annoying if your running other gunicorns on your machine which also get restarted
+                    # os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+                    # os.system("kill -1 `pgrep gunicorn`")
+    # todo refine
+    # if image != None:
+    #     image = refiner(
+    #         prompt=prompt,
+    #         # width=block_width,
+    #         # height=block_height,
+    #         num_inference_steps=n_steps,
+    #         # denoising_start=high_noise_frac,
+    #         image=image,
+    #     ).images[0]
+    if width != block_width or height != block_height:
+        # resize to original size width/height
+        # find aspect ratio to scale up to that covers the original img input width/height
+        scale_up_ratio = max(width / block_width, height / block_height)
+        image = image.resize((math.ceil(block_width * scale_up_ratio), math.ceil(height * scale_up_ratio)))
+        # crop image to original size
+        image = image.crop((0, 0, width, height))
+    # try:
+    #     # gc.collect()
+    #     torch.cuda.empty_cache()
+    # except Exception as e:
+    #     traceback.print_exc()
+    #     logger.info("restarting server to fix cuda issues (device side asserts)")
+    #     # todo fix device side asserts instead of restart to fix
+    #     # todo only restart the correct gunicorn
+    #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+    #     os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+    #     os.system("kill -1 `pgrep gunicorn`")
+    # save as bytesio
+    bs = BytesIO()
+
+    bright_count = np.sum(np.array(image) > 0)
+    if bright_count == 0:
+        # we have a black image, this is an error likely we need a restart
+        logger.info("restarting server to fix cuda issues (device side asserts)")
+        #     # todo fix device side asserts instead of restart to fix
+        #     # todo only restart the correct gunicorn
+        #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+        os.system("kill -1 `pgrep gunicorn`")
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep uvicorn`")
+        os.system("kill -1 `pgrep uvicorn`")
+
+        return None
+    image.save(bs, quality=85, optimize=True, format="webp")
+    bio = bs.getvalue()
+    # touch progress.txt file - if we dont do this we get restarted by supervisor/other processes for reliability
+    with open("progress.txt", "w") as f:
+        current_time = datetime.now().strftime("%H:%M:%S")
+        f.write(f"{current_time}")
+    return bio
+
+def inpaint_image_from_prompt(prompt, image_url: str, mask_url: str):
+    prompt = shorten_too_long_text(prompt)
+    # image = pipe(prompt=prompt).images[0]
+
+    init_image = load_image(image_url).convert("RGB")
+    mask_image = load_image(mask_url).convert("RGB") # why rgb for a 1 channel mask?
+    num_inference_steps = 75
+    high_noise_frac = 0.7
+
+    try:
+        image = inpaintpipe(
+            prompt=prompt,
+            image=init_image,
+            mask_image=mask_image,
+            num_inference_steps=num_inference_steps,
+            denoising_start=high_noise_frac,
+            output_type="latent",
+        ).images[0]  # normally uses 50 steps
+    except Exception as e:
+        # try rm stopwords + half the prompt
+        # todo try prompt permutations
+        logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+        prompt = ' '.join((word for word in prompt if word not in stopwords))
+        prompts = prompt.split()
+
+        prompt = ' '.join(prompts[:len(prompts) // 2])
+        logger.info(f"shortened prompt to: {len(prompt)}")
+        image = None
+        if prompt:
+            try:
+                image = pipe(
+                    prompt=prompt,
+                    image=init_image,
+                    mask_image=mask_image,
+                    num_inference_steps=num_inference_steps,
+                    denoising_start=high_noise_frac,
+                    output_type="latent",
+                ).images[0]  # normally uses 50 steps
+            except Exception as e:
+                # logger.info("trying to permute prompt")
+                # # try two swaps of the prompt/permutations
+                # prompt = prompt.split()
+                # prompt = ' '.join(permutations(prompt, 2).__next__())
+                logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+                prompt = ' '.join((word for word in prompt if word not in stopwords))
+                prompts = prompt.split()
+
+                prompt = ' '.join(prompts[:len(prompts) // 2])
+                logger.info(f"shortened prompt to: {len(prompt)}")
+
+                try:
+                    image = inpaintpipe(
+                        prompt=prompt,
+                        image=init_image,
+                        mask_image=mask_image,
+                        num_inference_steps=num_inference_steps,
+                        denoising_start=high_noise_frac,
+                        output_type="latent",
+                    ).images[0]  # normally uses 50 steps
+                except Exception as e:
+                    # just error out
+                    traceback.print_exc()
+                    raise e
+                    # logger.info("restarting server to fix cuda issues (device side asserts)")
+                    # todo fix device side asserts instead of restart to fix
+                    # todo only restart the correct gunicorn
+                    # this could be really annoying if your running other gunicorns on your machine which also get restarted
+                    # os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+                    # os.system("kill -1 `pgrep gunicorn`")
+    if image != None:
+        image = inpaint_refiner(
+            prompt=prompt,
+            image=image,
+            mask_image=mask_image,
+            num_inference_steps=num_inference_steps,
+            denoising_start=high_noise_frac,
+
+        ).images[0]
+    # try:
+    #     # gc.collect()
+    #     torch.cuda.empty_cache()
+    # except Exception as e:
+    #     traceback.print_exc()
+    #     logger.info("restarting server to fix cuda issues (device side asserts)")
+    #     # todo fix device side asserts instead of restart to fix
+    #     # todo only restart the correct gunicorn
+    #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+    #     os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+    #     os.system("kill -1 `pgrep gunicorn`")
+    # save as bytesio
+    bs = BytesIO()
+
+    bright_count = np.sum(np.array(image) > 0)
+    if bright_count == 0:
+        # we have a black image, this is an error likely we need a restart
+        logger.info("restarting server to fix cuda issues (device side asserts)")
+        #     # todo fix device side asserts instead of restart to fix
+        #     # todo only restart the correct gunicorn
+        #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+        os.system("kill -1 `pgrep gunicorn`")
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep uvicorn`")
+        os.system("kill -1 `pgrep uvicorn`")
+
+        return None
+    image.save(bs, quality=85, optimize=True, format="webp")
+    bio = bs.getvalue()
+    # touch progress.txt file - if we dont do this we get restarted by supervisor/other processes for reliability
+    with open("progress.txt", "w") as f:
+        current_time = datetime.now().strftime("%H:%M:%S")
+        f.write(f"{current_time}")
+    return bio
+
+
+
+def shorten_too_long_text(prompt):
+    if len(prompt) > 200:
+        # remove stopwords
+        prompt = prompt.split() # todo also split hyphens
+        prompt = ' '.join((word for word in prompt if word not in stopwords))
+        if len(prompt) > 200:
+            prompt = prompt[:200]
+    return prompt
+
+# image = pipe(prompt=prompt).images[0]
+#
+# image.save("test.png")
+# # save all images
+# for i, image in enumerate(images):
+#     image.save(f"{i}.png")

img/pr2/main.py

@@ -0,0 +1,528 @@
+import gc
+import math
+import multiprocessing
+import os
+import traceback
+from datetime import datetime
+from io import BytesIO
+from itertools import permutations
+from multiprocessing.pool import Pool
+from pathlib import Path
+from urllib.parse import quote_plus
+
+import numpy as np
+import nltk
+import torch
+
+from PIL.Image import Image
+from diffusers import DiffusionPipeline, StableDiffusionXLInpaintPipeline
+from diffusers.utils import load_image
+from fastapi import FastAPI
+from fastapi.middleware.gzip import GZipMiddleware
+from loguru import logger
+from starlette.middleware.cors import CORSMiddleware
+from starlette.responses import FileResponse
+from starlette.responses import JSONResponse
+
+from env import BUCKET_PATH, BUCKET_NAME
+# from stable_diffusion_server.bucket_api import check_if_blob_exists, upload_to_bucket
+torch._dynamo.config.suppress_errors = True
+
+import string
+import random
+
+def generate_save_path():
+    # initializing size of string
+    N = 7
+
+    # using random.choices()
+    # generating random strings
+    res = ''.join(random.choices(string.ascii_uppercase +
+                                 string.digits, k=N))
+    return res
+
+pipe = DiffusionPipeline.from_pretrained(
+    "models/stable-diffusion-xl-base-1.0",
+    torch_dtype=torch.bfloat16,
+    use_safetensors=True,
+    variant="fp16",
+    # safety_checker=None,
+)  # todo try torch_dtype=bfloat16
+pipe.watermark = None
+
+pipe.to("cuda")
+
+refiner = DiffusionPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-refiner-1.0",
+    text_encoder_2=pipe.text_encoder_2,
+    vae=pipe.vae,
+    torch_dtype=torch.bfloat16, # safer to use bfloat?
+    use_safetensors=True,
+    variant="fp16", #remember not to download the big model
+)
+refiner.watermark = None
+refiner.to("cuda")
+
+# {'scheduler', 'text_encoder', 'text_encoder_2', 'tokenizer', 'tokenizer_2', 'unet', 'vae'} can be passed in from existing model
+inpaintpipe = StableDiffusionXLInpaintPipeline.from_pretrained(
+    "models/stable-diffusion-xl-base-1.0", torch_dtype=torch.bfloat16, variant="fp16", use_safetensors=True,
+    scheduler=pipe.scheduler,
+    text_encoder=pipe.text_encoder,
+    text_encoder_2=pipe.text_encoder_2,
+    tokenizer=pipe.tokenizer,
+    tokenizer_2=pipe.tokenizer_2,
+    unet=pipe.unet,
+    vae=pipe.vae,
+    # load_connected_pipeline=
+)
+# # switch out to save gpu mem
+# del inpaintpipe.vae
+# del inpaintpipe.text_encoder_2
+# del inpaintpipe.text_encoder
+# del inpaintpipe.scheduler
+# del inpaintpipe.tokenizer
+# del inpaintpipe.tokenizer_2
+# del inpaintpipe.unet
+# inpaintpipe.vae = pipe.vae
+# inpaintpipe.text_encoder_2 = pipe.text_encoder_2
+# inpaintpipe.text_encoder = pipe.text_encoder
+# inpaintpipe.scheduler = pipe.scheduler
+# inpaintpipe.tokenizer = pipe.tokenizer
+# inpaintpipe.tokenizer_2 = pipe.tokenizer_2
+# inpaintpipe.unet = pipe.unet
+# todo this should work
+# inpaintpipe = StableDiffusionXLInpaintPipeline( # construct an inpainter using the existing model
+#     vae=pipe.vae,
+#     text_encoder_2=pipe.text_encoder_2,
+#     text_encoder=pipe.text_encoder,
+#     unet=pipe.unet,
+#     scheduler=pipe.scheduler,
+#     tokenizer=pipe.tokenizer,
+#     tokenizer_2=pipe.tokenizer_2,
+#     requires_aesthetics_score=False,
+# )
+inpaintpipe.to("cuda")
+inpaintpipe.watermark = None
+# inpaintpipe.register_to_config(requires_aesthetics_score=False)
+
+inpaint_refiner = StableDiffusionXLInpaintPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-refiner-1.0",
+    text_encoder_2=inpaintpipe.text_encoder_2,
+    vae=inpaintpipe.vae,
+    torch_dtype=torch.bfloat16,
+    use_safetensors=True,
+    variant="fp16",
+
+    tokenizer_2=refiner.tokenizer_2,
+    tokenizer=refiner.tokenizer,
+    scheduler=refiner.scheduler,
+    text_encoder=refiner.text_encoder,
+    unet=refiner.unet,
+)
+# del inpaint_refiner.vae
+# del inpaint_refiner.text_encoder_2
+# del inpaint_refiner.text_encoder
+# del inpaint_refiner.scheduler
+# del inpaint_refiner.tokenizer
+# del inpaint_refiner.tokenizer_2
+# del inpaint_refiner.unet
+# inpaint_refiner.vae = inpaintpipe.vae
+# inpaint_refiner.text_encoder_2 = inpaintpipe.text_encoder_2
+#
+# inpaint_refiner.text_encoder = refiner.text_encoder
+# inpaint_refiner.scheduler = refiner.scheduler
+# inpaint_refiner.tokenizer = refiner.tokenizer
+# inpaint_refiner.tokenizer_2 = refiner.tokenizer_2
+# inpaint_refiner.unet = refiner.unet
+
+# inpaint_refiner = StableDiffusionXLInpaintPipeline(
+#     text_encoder_2=inpaintpipe.text_encoder_2,
+#     vae=inpaintpipe.vae,
+#     # the rest from the existing refiner
+#     tokenizer_2=refiner.tokenizer_2,
+#     tokenizer=refiner.tokenizer,
+#     scheduler=refiner.scheduler,
+#     text_encoder=refiner.text_encoder,
+#     unet=refiner.unet,
+#     requires_aesthetics_score=False,
+# )
+inpaint_refiner.to("cuda")
+inpaint_refiner.watermark = None
+# inpaint_refiner.register_to_config(requires_aesthetics_score=False)
+
+n_steps = 40
+high_noise_frac = 0.8
+
+# if using torch < 2.0
+# pipe.enable_xformers_memory_efficient_attention()
+
+
+# pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
+# this can cause errors on some inputs so consider disabling it
+pipe.unet = torch.compile(pipe.unet)
+refiner.unet = torch.compile(refiner.unet)#, mode="reduce-overhead", fullgraph=True)
+# compile the inpainters - todo reuse the other unets? swap out the models for others/del them so they share models and can be swapped efficiently
+inpaintpipe.unet = pipe.unet
+inpaint_refiner.unet = refiner.unet
+# inpaintpipe.unet = torch.compile(inpaintpipe.unet)
+# inpaint_refiner.unet = torch.compile(inpaint_refiner.unet)
+from pydantic import BaseModel
+
+app = FastAPI(
+    openapi_url="/static/openapi.json",
+    docs_url="/swagger-docs",
+    redoc_url="/redoc",
+    title="Generate Images Netwrck API",
+    description="Character Chat API",
+    # root_path="https://api.text-generator.io",
+    version="1",
+)
+app.add_middleware(GZipMiddleware, minimum_size=1000)
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+stopwords = nltk.corpus.stopwords.words("english")
+
+class Img(BaseModel):
+    system_prompt: str
+    ASSISTANT: str
+
+# img_url = "http://phlrr2019.guest.corp.microsoft.com:8000/img1_sdv2.1.png"
+img_url = "http://phlrr3058.guest.corp.microsoft.com:8000/"#/img1_sdv2.1.png"
+
+@app.post("/image_url")
+def image_url(img: Img):
+    system_prompt = img.system_prompt
+    prompt = img.ASSISTANT
+    # if Path(save_path).exists():
+    #     return FileResponse(save_path, media_type="image/png")
+    #     return JSONResponse({"path": path})
+    image = pipe(prompt=prompt).images[0]
+    # if not save_path:
+    save_path = generate_save_path()
+    save_path = f"images/{save_path}.png"
+    image.save(save_path)
+    # save_path = '/'.join(path_components) + quote_plus(final_name)
+    path = f"{img_url}/{save_path}"
+    return JSONResponse({"path": path})
+
+
+@app.get("/make_image")
+# @app.post("/make_image")
+def make_image(prompt: str, save_path: str = ""):
+    if Path(save_path).exists():
+        return FileResponse(save_path, media_type="image/png")
+    image = pipe(prompt=prompt).images[0]
+    if not save_path:
+        save_path = f"images/{prompt}.png"
+    image.save(save_path)
+    return FileResponse(save_path, media_type="image/png")
+
+
+@app.get("/create_and_upload_image")
+def create_and_upload_image(prompt: str, width: int=1024, height:int=1024, save_path: str = ""):
+    path_components = save_path.split("/")[0:-1]
+    final_name = save_path.split("/")[-1]
+    if not path_components:
+        path_components = []
+    save_path = '/'.join(path_components) + quote_plus(final_name)
+    path = get_image_or_create_upload_to_cloud_storage(prompt, width, height, save_path)
+    return JSONResponse({"path": path})
+
+@app.get("/inpaint_and_upload_image")
+def inpaint_and_upload_image(prompt: str, image_url:str, mask_url:str, save_path: str = ""):
+    path_components = save_path.split("/")[0:-1]
+    final_name = save_path.split("/")[-1]
+    if not path_components:
+        path_components = []
+    save_path = '/'.join(path_components) + quote_plus(final_name)
+    path = get_image_or_inpaint_upload_to_cloud_storage(prompt, image_url, mask_url, save_path)
+    return JSONResponse({"path": path})
+
+
+def get_image_or_create_upload_to_cloud_storage(prompt:str,width:int, height:int, save_path:str):
+    prompt = shorten_too_long_text(prompt)
+    save_path = shorten_too_long_text(save_path)
+    # check exists - todo cache this
+    if check_if_blob_exists(save_path):
+        return f"https://{BUCKET_NAME}/{BUCKET_PATH}/{save_path}"
+    bio = create_image_from_prompt(prompt, width, height)
+    if bio is None:
+        return None # error thrown in pool
+    link = upload_to_bucket(save_path, bio, is_bytesio=True)
+    return link
+def get_image_or_inpaint_upload_to_cloud_storage(prompt:str, image_url:str, mask_url:str, save_path:str):
+    prompt = shorten_too_long_text(prompt)
+    save_path = shorten_too_long_text(save_path)
+    # check exists - todo cache this
+    if check_if_blob_exists(save_path):
+        return f"https://{BUCKET_NAME}/{BUCKET_PATH}/{save_path}"
+    bio = inpaint_image_from_prompt(prompt, image_url, mask_url)
+    if bio is None:
+        return None # error thrown in pool
+    link = upload_to_bucket(save_path, bio, is_bytesio=True)
+    return link
+
+# multiprocessing.set_start_method('spawn', True)
+# processes_pool = Pool(1) # cant do too much at once or OOM errors happen
+# def create_image_from_prompt_sync(prompt):
+#     """have to call this sync to avoid OOM errors"""
+#     return processes_pool.apply_async(create_image_from_prompt, args=(prompt,), ).wait()
+
+def create_image_from_prompt(prompt, width, height):
+    # round width and height down to multiple of 64
+    block_width = width - (width % 64)
+    block_height = height - (height % 64)
+    prompt = shorten_too_long_text(prompt)
+    # image = pipe(prompt=prompt).images[0]
+    try:
+        image = pipe(prompt=prompt,
+                     width=block_width,
+                     height=block_height,
+                     # denoising_end=high_noise_frac,
+                     # output_type='latent',
+                     # height=512,
+                     # width=512,
+                     num_inference_steps=50).images[0]  # normally uses 50 steps
+    except Exception as e:
+        # try rm stopwords + half the prompt
+        # todo try prompt permutations
+        logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+        prompt = ' '.join((word for word in prompt if word not in stopwords))
+        prompts = prompt.split()
+
+        prompt = ' '.join(prompts[:len(prompts) // 2])
+        logger.info(f"shortened prompt to: {len(prompt)}")
+        image = None
+        if prompt:
+            try:
+                image = pipe(prompt=prompt,
+                             width=block_width,
+                             height=block_height,
+                             # denoising_end=high_noise_frac,
+                             # output_type='latent',
+                             # height=512,
+                             # width=512,
+                             num_inference_steps=50).images[0]  # normally uses 50 steps
+            except Exception as e:
+                # logger.info("trying to permute prompt")
+                # # try two swaps of the prompt/permutations
+                # prompt = prompt.split()
+                # prompt = ' '.join(permutations(prompt, 2).__next__())
+                logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+                prompt = ' '.join((word for word in prompt if word not in stopwords))
+                prompts = prompt.split()
+
+                prompt = ' '.join(prompts[:len(prompts) // 2])
+                logger.info(f"shortened prompt to: {len(prompt)}")
+
+                try:
+                    image = pipe(prompt=prompt,
+                                 width=block_width,
+                                 height=block_height,
+                                 # denoising_end=high_noise_frac,
+                                 # output_type='latent', # dont need latent yet - we refine the image at full res
+                                 # height=512,
+                                 # width=512,
+                                 num_inference_steps=50).images[0]  # normally uses 50 steps
+                except Exception as e:
+                    # just error out
+                    traceback.print_exc()
+                    raise e
+                    # logger.info("restarting server to fix cuda issues (device side asserts)")
+                    # todo fix device side asserts instead of restart to fix
+                    # todo only restart the correct gunicorn
+                    # this could be really annoying if your running other gunicorns on your machine which also get restarted
+                    # os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+                    # os.system("kill -1 `pgrep gunicorn`")
+    # todo refine
+    # if image != None:
+    #     image = refiner(
+    #         prompt=prompt,
+    #         # width=block_width,
+    #         # height=block_height,
+    #         num_inference_steps=n_steps,
+    #         # denoising_start=high_noise_frac,
+    #         image=image,
+    #     ).images[0]
+    if width != block_width or height != block_height:
+        # resize to original size width/height
+        # find aspect ratio to scale up to that covers the original img input width/height
+        scale_up_ratio = max(width / block_width, height / block_height)
+        image = image.resize((math.ceil(block_width * scale_up_ratio), math.ceil(height * scale_up_ratio)))
+        # crop image to original size
+        image = image.crop((0, 0, width, height))
+    # try:
+    #     # gc.collect()
+    #     torch.cuda.empty_cache()
+    # except Exception as e:
+    #     traceback.print_exc()
+    #     logger.info("restarting server to fix cuda issues (device side asserts)")
+    #     # todo fix device side asserts instead of restart to fix
+    #     # todo only restart the correct gunicorn
+    #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+    #     os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+    #     os.system("kill -1 `pgrep gunicorn`")
+    # save as bytesio
+    bs = BytesIO()
+
+    bright_count = np.sum(np.array(image) > 0)
+    if bright_count == 0:
+        # we have a black image, this is an error likely we need a restart
+        logger.info("restarting server to fix cuda issues (device side asserts)")
+        #     # todo fix device side asserts instead of restart to fix
+        #     # todo only restart the correct gunicorn
+        #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+        os.system("kill -1 `pgrep gunicorn`")
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep uvicorn`")
+        os.system("kill -1 `pgrep uvicorn`")
+
+        return None
+    image.save(bs, quality=85, optimize=True, format="webp")
+    bio = bs.getvalue()
+    # touch progress.txt file - if we dont do this we get restarted by supervisor/other processes for reliability
+    with open("progress.txt", "w") as f:
+        current_time = datetime.now().strftime("%H:%M:%S")
+        f.write(f"{current_time}")
+    return bio
+
+def inpaint_image_from_prompt(prompt, image_url: str, mask_url: str):
+    prompt = shorten_too_long_text(prompt)
+    # image = pipe(prompt=prompt).images[0]
+
+    init_image = load_image(image_url).convert("RGB")
+    mask_image = load_image(mask_url).convert("RGB") # why rgb for a 1 channel mask?
+    num_inference_steps = 75
+    high_noise_frac = 0.7
+
+    try:
+        image = inpaintpipe(
+            prompt=prompt,
+            image=init_image,
+            mask_image=mask_image,
+            num_inference_steps=num_inference_steps,
+            denoising_start=high_noise_frac,
+            output_type="latent",
+        ).images[0]  # normally uses 50 steps
+    except Exception as e:
+        # try rm stopwords + half the prompt
+        # todo try prompt permutations
+        logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+        prompt = ' '.join((word for word in prompt if word not in stopwords))
+        prompts = prompt.split()
+
+        prompt = ' '.join(prompts[:len(prompts) // 2])
+        logger.info(f"shortened prompt to: {len(prompt)}")
+        image = None
+        if prompt:
+            try:
+                image = pipe(
+                    prompt=prompt,
+                    image=init_image,
+                    mask_image=mask_image,
+                    num_inference_steps=num_inference_steps,
+                    denoising_start=high_noise_frac,
+                    output_type="latent",
+                ).images[0]  # normally uses 50 steps
+            except Exception as e:
+                # logger.info("trying to permute prompt")
+                # # try two swaps of the prompt/permutations
+                # prompt = prompt.split()
+                # prompt = ' '.join(permutations(prompt, 2).__next__())
+                logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+                prompt = ' '.join((word for word in prompt if word not in stopwords))
+                prompts = prompt.split()
+
+                prompt = ' '.join(prompts[:len(prompts) // 2])
+                logger.info(f"shortened prompt to: {len(prompt)}")
+
+                try:
+                    image = inpaintpipe(
+                        prompt=prompt,
+                        image=init_image,
+                        mask_image=mask_image,
+                        num_inference_steps=num_inference_steps,
+                        denoising_start=high_noise_frac,
+                        output_type="latent",
+                    ).images[0]  # normally uses 50 steps
+                except Exception as e:
+                    # just error out
+                    traceback.print_exc()
+                    raise e
+                    # logger.info("restarting server to fix cuda issues (device side asserts)")
+                    # todo fix device side asserts instead of restart to fix
+                    # todo only restart the correct gunicorn
+                    # this could be really annoying if your running other gunicorns on your machine which also get restarted
+                    # os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+                    # os.system("kill -1 `pgrep gunicorn`")
+    if image != None:
+        image = inpaint_refiner(
+            prompt=prompt,
+            image=image,
+            mask_image=mask_image,
+            num_inference_steps=num_inference_steps,
+            denoising_start=high_noise_frac,
+
+        ).images[0]
+    # try:
+    #     # gc.collect()
+    #     torch.cuda.empty_cache()
+    # except Exception as e:
+    #     traceback.print_exc()
+    #     logger.info("restarting server to fix cuda issues (device side asserts)")
+    #     # todo fix device side asserts instead of restart to fix
+    #     # todo only restart the correct gunicorn
+    #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+    #     os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+    #     os.system("kill -1 `pgrep gunicorn`")
+    # save as bytesio
+    bs = BytesIO()
+
+    bright_count = np.sum(np.array(image) > 0)
+    if bright_count == 0:
+        # we have a black image, this is an error likely we need a restart
+        logger.info("restarting server to fix cuda issues (device side asserts)")
+        #     # todo fix device side asserts instead of restart to fix
+        #     # todo only restart the correct gunicorn
+        #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+        os.system("kill -1 `pgrep gunicorn`")
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep uvicorn`")
+        os.system("kill -1 `pgrep uvicorn`")
+
+        return None
+    image.save(bs, quality=85, optimize=True, format="webp")
+    bio = bs.getvalue()
+    # touch progress.txt file - if we dont do this we get restarted by supervisor/other processes for reliability
+    with open("progress.txt", "w") as f:
+        current_time = datetime.now().strftime("%H:%M:%S")
+        f.write(f"{current_time}")
+    return bio
+
+
+
+def shorten_too_long_text(prompt):
+    if len(prompt) > 200:
+        # remove stopwords
+        prompt = prompt.split() # todo also split hyphens
+        prompt = ' '.join((word for word in prompt if word not in stopwords))
+        if len(prompt) > 200:
+            prompt = prompt[:200]
+    return prompt
+
+# image = pipe(prompt=prompt).images[0]
+#
+# image.save("test.png")
+# # save all images
+# for i, image in enumerate(images):
+#     image.save(f"{i}.png")

img/readme.md

@@ -0,0 +1,109 @@
+simple stable diffusion server that saves images to cloud storage - returns links to google cloud storage
+
+## Creators
+[![netwrck logo](https://static.netwrck.com/static/img/netwrck-logo-colord256.png)](https://netwrck.com)
+
+Checkout [Voiced AI Characters to chat with](https://netwrck.com) at [netwrck.com](https://netwrck.com)
+
+Characters are narrated and written by many GPT models trained on 1000s of fantasy novels and chats.
+
+Also for LLMs for making Text - Checkout [Text-Generator.io](https://text-generator.io) for a Open Source text generator that uses many AI models to generate the best along with image understanding and OCR networks.
+## Setup
+
+. Create a virtual environment (optional)
+
+```bash
+python3 -m venv venv
+source venv/bin/activate
+```
+
+#### Install dependencies
+
+```bash
+pip install -r requirements.txt
+pip install -r dev-requirements.txt
+
+cd models
+git clone https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0    
+
+# install stopwords
+python -c "import nltk; nltk.download('stopwords')"
+```
+
+#### Edit settings in env.py
+#### download your Google cloud credentials to secrets/google-credentials.json
+Images generated will be stored in your bucket
+#### Run the server
+
+```bash
+GOOGLE_APPLICATION_CREDENTIALS=secrets/google-credentials.json gunicorn  -k uvicorn.workers.UvicornWorker -b :8000 main:app --timeout 600 -w 1 
+```
+
+with max 4 requests at a time
+This will drop a lot of requests under load instead of taking on too much work and causing OOM Errors.
+
+```bash
+GOOGLE_APPLICATION_CREDENTIALS=secrets/google-credentials.json PYTHONPATH=. uvicorn --port 8000 --timeout-keep-alive 600 --workers 1 --backlog 1 --limit-concurrency 4 main:app
+```
+
+#### Make a Request
+
+http://localhost:8000/create_and_upload_image?prompt=good%20looking%20elf%20fantasy%20character&save_path=created/elf.webp
+
+Response
+```shell
+{"path":"https://storage.googleapis.com/static.netwrck.com/static/uploads/created/elf.png"}
+```
+
+http://localhost:8000/docs
+
+
+Check to see that "good Looking elf fantasy character" was created
+
+![elf.png](https://storage.googleapis.com/static.netwrck.com/static/uploads/created/elf.png)
+![elf2.png](https://storage.googleapis.com/static.netwrck.com/static/uploads/created/elf2.png)
+
+### Testing
+
+```bash
+GOOGLE_APPLICATION_CREDENTIALS=secrets/google-credentials.json pytest .
+```
+
+
+#### Running under supervisord
+
+edit ops/supervisor.conf
+
+install the supervisor
+apt-get install -y supervisor
+```bash
+sudo cat >/etc/supervisor/conf.d/python-app.conf << EOF
+[program:sdif_http_server]
+directory=/home/lee/code/sdif
+command=/home/lee/code/sdif/.env/bin/uvicorn --port 8000 --timeout-keep-alive 600 --workers 1 --backlog 1 --limit-concurrency 4 main:app
+autostart=true
+autorestart=true
+environment=VIRTUAL_ENV="/home/lee/code/sdif/.env/",PATH="/opt/app/sdif/.env/bin",HOME="/home/lee",GOOGLE_APPLICATION_CREDENTIALS="secrets/google-credentials.json",PYTHONPATH="/home/lee/code/sdif"
+stdout_logfile=syslog
+stderr_logfile=syslog
+user=lee
+EOF
+
+supervisorctl reread
+supervisorctl update
+```
+
+#### run a manager process to kill/restart if the server if it is hanging
+
+Sometimes the server just stops working and needs a hard restart
+
+This command will kill the server if it is hanging and restart it (must be running under supervisorctl)
+```
+python3 manager.py
+```
+
+# hack restarting without supervisor
+run the server in a infinite loop
+```
+while true; do GOOGLE_APPLICATION_CREDENTIALS=secrets/google-credentials.json PYTHONPATH=. uvicorn --port 8000 --timeout-keep-alive 600 --workers 1 --backlog 1 --limit-concurrency 4 main:app; done
+```

img/requirements.txt

@@ -0,0 +1,67 @@
+accelerate==0.20.3
+annotated-types==0.5.0
+anyio==3.7.1
+certifi==2023.5.7
+charset-normalizer==3.2.0
+click==8.1.4
+cmake==3.26.4
+diffusers==0.20.0
+exceptiongroup==1.1.2
+fastapi==0.100.0
+filelock==3.12.2
+fsspec==2023.6.0
+gunicorn==20.1.0
+h11==0.14.0
+huggingface-hub==0.16.4
+idna==3.4
+importlib-metadata==6.8.0
+invisible-watermark==0.2.0
+Jinja2==3.1.2
+lit==16.0.6
+MarkupSafe==2.1.3
+mpmath==1.3.0
+networkx==3.1
+numpy==1.25.0
+opencv-python==4.8.0.74
+packaging==23.1
+Pillow==10.0.0
+psutil==5.9.5
+pydantic==2.0.2
+pydantic_core==2.1.2
+PyWavelets==1.4.1
+PyYAML==6.0
+regex==2023.6.3
+requests==2.31.0
+safetensors==0.3.1
+sniffio==1.3.0
+starlette==0.27.0
+sympy==1.12
+tokenizers==0.13.3
+torch==2.0.1
+tqdm==4.65.0
+transformers==4.30.2
+#triton==2.0.0
+typing_extensions==4.7.1
+urllib3==2.0.3
+uvicorn==0.22.0
+zipp==3.15.0
+jinja2
+loguru==0.6.0
+
+google-api-python-client==2.43.0
+google-api-core #1.31.5
+#google-cloud-storage==2.3.0 #not on gae python
+google-cloud-storage==2.0.0
+
+google-cloud-ndb==1.11.1
+cachetools==4.2.4
+
+python-multipart==0.0.6
+nltk==3.8.1
+diskcache==5.5.1
+
+protobuf==3.19.5
+google-cloud-aiplatform==1.25.0
+# openai==0.27.7
+# requests==2.28.2
+# rollbar==0.16.3

img/scripts/test_compression.py

@@ -0,0 +1,22 @@
+# save images in 1-10 compresion timing the results
+from pathlib import Path
+from time import time
+def test_compression():
+    save_dir = Path("./imgs-sd/test/")
+    save_dir.mkdir(exist_ok=True, parents=True)
+
+    from PIL import Image
+
+    image = Image.open("/home/lee/code/sdif/imgs-sd/Woody.png").convert("RGB")
+    start = time()
+
+    image.save(save_dir / f"woody-.webp", format="webp")
+    end = time()
+    print(f"Time to save image with quality : {end - start}")
+
+    for i in range(0, 100):
+        start = time()
+
+        image.save(save_dir / f"woody-{i}.webp", quality=i, optimize=True, format="webp")
+        end = time()
+        print(f"Time to save image with quality {i}: {end - start}")

img/stable-diffusion-server/.gitignore

@@ -0,0 +1,13 @@
+models
+sd-images1
+imgs-sd
+images
+backdrops
+.env
+venv
+secrets
+.pytest_cache
+progress.txt
+.idea
+__pycache__
+

img/stable-diffusion-server/dev-requirements.txt

@@ -0,0 +1,11 @@
+pytest
+
+pytest-asyncio
+requests-futures==1.0.0
+httpx
+djlint
+pytest-env==0.8.1
+ipython
+
+line-profiler-pycharm==1.1.0
+line-profiler==4.0.3

img/stable-diffusion-server/env.py

@@ -0,0 +1,2 @@
+BUCKET_NAME = 'static.netwrck.com'
+BUCKET_PATH = 'static/uploads'

img/stable-diffusion-server/img2img.py

@@ -0,0 +1,25 @@
+import requests
+import torch
+from PIL import Image
+from io import BytesIO
+
+from diffusers import StableDiffusionImg2ImgPipeline
+
+device = "cuda"
+model_id_or_path = "runwayml/stable-diffusion-v1-5"
+# model_id_or_path = "models/stable-diffusion-xl-base-0.9"
+pipe = StableDiffusionImg2ImgPipeline.from_pretrained(model_id_or_path, torch_dtype=torch.float16, variant="fp16", safety_checker=None)
+pipe = pipe.to(device)
+
+url = "https://raw.githubusercontent.com/CompVis/stable-diffusion/main/assets/stable-samples/img2img/sketch-mountains-input.jpg"
+
+response = requests.get(url)
+# init_image = Image.open(BytesIO(response.content)).convert("RGB")
+init_image = Image.open("/mnt/c/Users/leepenkman/Pictures/aiknight-neon-punk-fantasy-art-good-looking-trending-fantastic-1.webp").convert("RGB")
+# init_image = init_image.resize((768, 512))
+init_image = init_image.resize((1920, 1080))
+
+prompt = "knight neon punk fantasy art good looking trending fantastic"
+
+images = pipe(prompt=prompt, image=init_image, strength=0.75, guidance_scale=7.5).images
+images[0].save("fantasy_landscape.png")

img/stable-diffusion-server/img2imgsd.py

@@ -0,0 +1,74 @@
+from pathlib import Path
+
+import numpy as np
+import requests
+import torch
+from PIL import Image
+from io import BytesIO
+
+# from diffusers import StableDiffusionImg2ImgPipeline
+
+# device = "cuda"
+# model_id_or_path = "runwayml/stable-diffusion-v1-5"
+# # model_id_or_path = "models/stable-diffusion-xl-base-0.9"
+# pipe = StableDiffusionImg2ImgPipeline.from_pretrained(model_id_or_path, torch_dtype=torch.float16, variant="fp16", safety_checker=None)
+# pipe = pipe.to(device)
+
+from diffusers import StableDiffusionXLImg2ImgPipeline
+from diffusers.utils import load_image
+
+from stable_diffusion_server.utils import log_time
+
+pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained(
+"stabilityai/stable-diffusion-xl-refiner-1.0",
+# "models/stable-diffusion-xl-base-0.9",
+    torch_dtype = torch.float16,
+    use_safetensors=True,
+    variant="fp16",
+)
+pipe = pipe.to("cuda") #  # "LayerNormKernelImpl" not implemented for 'Half' error if its on cpu it cant do fp16
+# idea composite: and re prompt img-img to support different sizes
+
+# url = "https://raw.githubusercontent.com/CompVis/stable-diffusion/main/assets/stable-samples/img2img/sketch-mountains-input.jpg"
+#
+# response = requests.get(url)
+# init_image = Image.open(BytesIO(response.content)).convert("RGB")
+# init_image = init_image.resize((768, 512))
+# successfully inpaints a deleted area strength=0.75
+# init_image = Image.open("/mnt/c/Users/leepenkman/Pictures/aiart/ainostalgic-colorful-relaxing-chill-realistic-cartoon-Charcoal-illustration-fantasy-fauvist-abstract-impressionist-watercolor-painting-Background-location-scenery-amazing-wonderful-Dog-Shelter-Worker-Dog.webp").convert("RGB")
+# redo something? strength 1
+# init_image = Image.open("/home/lee/code/sdif/mask.png").convert("RGB")
+init_image = Image.open("/mnt/c/Users/leepenkman/Pictures/dogstretch.png").convert("RGB")
+# init_image = Image.open("/mnt/c/Users/leepenkman/Pictures/dogcenter.png").convert("RGB")
+
+# init_image = init_image.resize((1080, 1920))
+init_image = init_image.resize((1920, 1080))
+# init_image = init_image.resize((1024, 1024))
+
+prompt = "A fantasy landscape, trending on artstation, beautiful amazing unreal surreal gorgeous impressionism"
+prompt = "mouth open nostalgic colorful relaxing chill realistic cartoon Charcoal illustration fantasy fauvist abstract impressionist watercolor painting Background location scenery amazing wonderful Dog Shelter Worker Dog"
+
+# images = pipe(prompt=prompt, image=init_image, strength=0.75, guidance_scale=7.5).images
+# images[0].save("fantasy_landscape.png")
+#
+# # url = "https://huggingface.co/datasets/patrickvonplaten/images/resolve/main/aa_xl/000000009.png"
+#
+# init_image = load_image(url).convert("RGB")
+# prompt = "a photo of an astronaut riding a horse on mars"
+study_dir = "images/study2"
+Path(study_dir).mkdir(parents=True, exist_ok=True)
+
+with log_time("img2img"):
+    with torch.inference_mode():
+        # for strength in range(.1, 1, .1):
+        for strength in np.linspace(.1, 1, 10):
+            image = pipe(prompt=prompt, image=init_image, strength=strength, guidance_scale=7.6).images[0]
+            image.save(
+                study_dir + "/fantasy_dogimgimgdogstretchopening" + str(strength) + "guidance_scale" + str(7.6) + ".png")
+        #     # for guidance_scale in range(1, 10, .5):
+        #     for guidance_scale in np.linspace(1, 100, 10):
+        #         image = pipe(prompt=prompt, image=init_image, strength=strength, guidance_scale=guidance_scale).images[0]
+        #         image.save("images/study/fantasy_dogimgimgdogstretch" + str(strength) + "guidance_scale" + str(guidance_scale) + ".png")
+        # image = pipe(prompt, image=init_image, strength=0.2, guidance_scale=7.5).images[0]
+        # image.save("images/fantasy_dogimgimgdogstretch.png")
+        # image.save("images/fantasy_dogimgimgdogcenter.png")

img/stable-diffusion-server/img2imgsdr.py

@@ -0,0 +1,53 @@
+import PIL.Image
+
+from diffusers import DiffusionPipeline
+import torch
+
+import numpy as np
+
+from stable_diffusion_server.utils import log_time
+
+pipe = DiffusionPipeline.from_pretrained(
+    "models/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, variant="fp16", use_safetensors=True
+)
+pipe.to("cuda")
+
+refiner = DiffusionPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-refiner-1.0",
+    text_encoder_2=pipe.text_encoder_2,
+    vae=pipe.vae,
+    torch_dtype=torch.float16,
+    use_safetensors=True,
+    variant="fp16",
+)
+refiner.to("cuda")
+
+prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
+use_refiner = True
+with log_time('diffuse'):
+    with torch.inference_mode():
+        image = pipe(prompt=prompt, output_type="latent" if use_refiner else "pil").images[0]
+        # experiment try deleting a whole bunch of pixels and see if the refiner can recreate them
+        # delete top 30% of pixels
+        # image = image[0:0.7]
+        #pixels to delete
+        # pixels_to_delete = int(0.3 * 1024)
+        # delete top 30% of pixels
+        # image.save("latent.png")
+        # image_data = PIL.Image.fromarray(image)
+        # image_data.save("latent.png")
+
+        # image = np.array(image)
+        pixels_to_delete = int(0.3 * image.shape[0])
+        idx_to_delete = np.ones(image.shape[0], dtype=bool, device="cuda")
+        idx_to_delete[:pixels_to_delete] = False
+        image[idx_to_delete] = [0,0,0]
+
+        # image_data = PIL.Image.fromarray(image)
+        # image_data.save("latentcleared.png")
+
+
+        image = refiner(prompt=prompt, image=image[None, :]).images[0]
+
+
+

img/stable-diffusion-server/inpaint.py

@@ -0,0 +1,62 @@
+import torch
+
+from diffusers import StableDiffusionXLInpaintPipeline
+from diffusers.utils import load_image
+
+from stable_diffusion_server.utils import log_time
+
+import numpy as np
+import PIL.Image
+
+pipe = StableDiffusionXLInpaintPipeline.from_pretrained(
+    "models/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, variant="fp16", use_safetensors=True
+)
+pipe.to("cuda")
+
+refiner = StableDiffusionXLInpaintPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-refiner-1.0",
+    text_encoder_2=pipe.text_encoder_2,
+    vae=pipe.vae,
+    torch_dtype=torch.float16,
+    use_safetensors=True,
+    variant="fp16",
+)
+refiner.to("cuda")
+
+img_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo.png"
+mask_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo_mask.png"
+# inpaint_and_upload_image?prompt=majestic tiger sitting on a bench&image_url=https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo.png&mask_url=https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo_mask.png&save_path=tests/inpaint.webp
+# inpainting can be used to upscale to 1080p
+
+
+init_image = load_image(img_url).convert("RGB")
+# mask_image = load_image(mask_url).convert("RGB")
+# mask image all ones same shape as init_image
+
+# here's a failed experiment: inpainting cannot be used as style transfer/it doesnt recreate ain image doing a full mask in this way
+image_size = init_image.size
+ones_of_size = np.ones(image_size, np.uint8) * 255
+mask_image = PIL.Image.fromarray(ones_of_size.astype(np.uint8))
+# mask_image = torch.ones_like(init_image) * 255
+prompt = "A majestic tiger sitting on a bench, castle backdrop elegent anime"
+num_inference_steps = 75
+high_noise_frac = 0.7
+with log_time("inpaint"):
+    with torch.inference_mode():
+        image = pipe(
+            prompt=prompt,
+            image=init_image,
+            mask_image=mask_image,
+            num_inference_steps=num_inference_steps,
+            denoising_start=high_noise_frac,
+            output_type="latent",
+        ).images
+        image = refiner(
+            prompt=prompt,
+            image=image,
+            mask_image=mask_image,
+            num_inference_steps=num_inference_steps,
+            denoising_start=high_noise_frac,
+        ).images[0]
+
+image.save("inpaintfull.png")

img/stable-diffusion-server/log.0925

@@ -0,0 +1,53 @@
+v-haipe+   551 16041 99 08:16 pts/2    00:00:17 python LiLa/gsm8k_cluster.py
+v-haipe+  9211 10235  3 Sep24 pts/10   00:32:12 python LiLa/chatgpt_evol_lila_gsm8k_domain.py --start 0 --end 2000
+v-haipe+  9288 10459  3 Sep24 pts/11   00:28:30 python LiLa/chatgpt_evol_lila_gsm8k_domain.py --start 2000 --end 4000
+v-haipe+  9310 10667  3 Sep24 pts/12   00:27:45 python LiLa/chatgpt_evol_lila_gsm8k_domain.py --start 4000 --end 6000
+v-haipe+  9341 10865  3 Sep24 pts/13   00:26:50 python LiLa/chatgpt_evol_lila_gsm8k_domain.py --start 6000 --end 8000
+v-haipe+  9379 25248  3 Sep24 pts/16   00:27:01 python LiLa/chatgpt_evol_lila_gsm8k_domain.py --start 8000 --end 10000
+v-haipe+  9410 25467  3 Sep24 pts/17   00:27:17 python LiLa/chatgpt_evol_lila_gsm8k_domain.py --start 10000 --end 12000
+v-haipe+  9438 26561  3 Sep24 pts/19   00:27:17 python LiLa/chatgpt_evol_lila_gsm8k_domain.py --start 12000 --end 14000
+v-haipe+  9469 26761  3 Sep24 pts/20   00:26:55 python LiLa/chatgpt_evol_lila_gsm8k_domain.py --start 14000 --end 16000
+v-haipe+  9500 26968  3 Sep24 pts/21   00:27:09 python LiLa/chatgpt_evol_lila_gsm8k_domain.py --start 16000 --end 18000
+v-haipe+  9531 27172  3 Sep24 pts/22   00:29:29 python LiLa/chatgpt_evol_lila_gsm8k_domain.py --start 18000 --end 20000
+v-haipe+  9775  9560  3 Sep24 pts/29   00:30:29 python LiLa/chatgpt_evol_lila_gsm8k_domain.py --start 20000 --end 22000
+v-haipe+ 11262 24577  0 Sep23 pts/8    00:00:06 python app.py
+v-haipe+ 11300 11262  0 Sep23 pts/8    00:20:54 /home/v-haipengluo/.conda/envs/wizardweb/bin/python /workspaceblobstore/qins/test/20220316/kai/research/code_repo/wizard_verse/code_repo/server_code/wizard_verse/lm/server_lm/app.py
+v-haipe+ 11604 20782 98 Sep23 pts/4    2-00:06:57 python -m vllm.entrypoints.api_server --model /workspaceblobstore/caxu/trained_models/13Bv2_497kcontinueroleplay_dsys_2048_e4_2e_5/checkpoint-75 --host phlrr3006.guest.corp.microsoft.com --port 7991
+v-haipe+ 13722 22601  0 Sep24 pts/6    00:09:37 /home/v-haipengluo/.conda/envs/sdxl/bin/python /home/v-haipengluo/.conda/envs/sdxl/bin/uvicorn --host=phlrr3006.guest.corp.microsoft.com --port 7999 --workers 1 --backlog 1 --limit-concurrency 4 main_v3:app
+v-haipe+ 13830 13722  0 Sep24 pts/6    00:00:05 /home/v-haipengluo/.conda/envs/sdxl/bin/python /home/v-haipengluo/.conda/envs/sdxl/bin/uvicorn --host=phlrr3006.guest.corp.microsoft.com --port 7999 --workers 1 --backlog 1 --limit-concurrency 4 main_v3:app
+v-haipe+ 13834 13722  0 Sep24 pts/6    00:00:05 /home/v-haipengluo/.conda/envs/sdxl/bin/python /home/v-haipengluo/.conda/envs/sdxl/bin/uvicorn --host=phlrr3006.guest.corp.microsoft.com --port 7999 --workers 1 --backlog 1 --limit-concurrency 4 main_v3:app
+v-haipe+ 13837 13722  0 Sep24 pts/6    00:00:05 /home/v-haipengluo/.conda/envs/sdxl/bin/python /home/v-haipengluo/.conda/envs/sdxl/bin/uvicorn --host=phlrr3006.guest.corp.microsoft.com --port 7999 --workers 1 --backlog 1 --limit-concurrency 4 main_v3:app
+v-haipe+ 13839 13722  0 Sep24 pts/6    00:00:05 /home/v-haipengluo/.conda/envs/sdxl/bin/python /home/v-haipengluo/.conda/envs/sdxl/bin/uvicorn --host=phlrr3006.guest.corp.microsoft.com --port 7999 --workers 1 --backlog 1 --limit-concurrency 4 main_v3:app
+v-haipe+ 13841 13722  0 Sep24 pts/6    00:00:05 /home/v-haipengluo/.conda/envs/sdxl/bin/python /home/v-haipengluo/.conda/envs/sdxl/bin/uvicorn --host=phlrr3006.guest.corp.microsoft.com --port 7999 --workers 1 --backlog 1 --limit-concurrency 4 main_v3:app
+v-haipe+ 13843 13722  0 Sep24 pts/6    00:00:05 /home/v-haipengluo/.conda/envs/sdxl/bin/python /home/v-haipengluo/.conda/envs/sdxl/bin/uvicorn --host=phlrr3006.guest.corp.microsoft.com --port 7999 --workers 1 --backlog 1 --limit-concurrency 4 main_v3:app
+v-haipe+ 13845 13722  0 Sep24 pts/6    00:00:05 /home/v-haipengluo/.conda/envs/sdxl/bin/python /home/v-haipengluo/.conda/envs/sdxl/bin/uvicorn --host=phlrr3006.guest.corp.microsoft.com --port 7999 --workers 1 --backlog 1 --limit-concurrency 4 main_v3:app
+v-haipe+ 13847 13722  0 Sep24 pts/6    00:00:05 /home/v-haipengluo/.conda/envs/sdxl/bin/python /home/v-haipengluo/.conda/envs/sdxl/bin/uvicorn --host=phlrr3006.guest.corp.microsoft.com --port 7999 --workers 1 --backlog 1 --limit-concurrency 4 main_v3:app
+v-haipe+ 13849 13722  0 Sep24 pts/6    00:00:05 /home/v-haipengluo/.conda/envs/sdxl/bin/python /home/v-haipengluo/.conda/envs/sdxl/bin/uvicorn --host=phlrr3006.guest.corp.microsoft.com --port 7999 --workers 1 --backlog 1 --limit-concurrency 4 main_v3:app
+v-haipe+ 13851 13722  0 Sep24 pts/6    00:00:05 /home/v-haipengluo/.conda/envs/sdxl/bin/python /home/v-haipengluo/.conda/envs/sdxl/bin/uvicorn --host=phlrr3006.guest.corp.microsoft.com --port 7999 --workers 1 --backlog 1 --limit-concurrency 4 main_v3:app
+v-haipe+ 13853 13722  0 Sep24 pts/6    00:00:05 /home/v-haipengluo/.conda/envs/sdxl/bin/python /home/v-haipengluo/.conda/envs/sdxl/bin/uvicorn --host=phlrr3006.guest.corp.microsoft.com --port 7999 --workers 1 --backlog 1 --limit-concurrency 4 main_v3:app
+v-haipe+ 13855 13722  0 Sep24 pts/6    00:00:05 /home/v-haipengluo/.conda/envs/sdxl/bin/python /home/v-haipengluo/.conda/envs/sdxl/bin/uvicorn --host=phlrr3006.guest.corp.microsoft.com --port 7999 --workers 1 --backlog 1 --limit-concurrency 4 main_v3:app
+v-haipe+ 13857 13722  0 Sep24 pts/6    00:00:05 /home/v-haipengluo/.conda/envs/sdxl/bin/python /home/v-haipengluo/.conda/envs/sdxl/bin/uvicorn --host=phlrr3006.guest.corp.microsoft.com --port 7999 --workers 1 --backlog 1 --limit-concurrency 4 main_v3:app
+v-haipe+ 13859 13722  0 Sep24 pts/6    00:00:05 /home/v-haipengluo/.conda/envs/sdxl/bin/python /home/v-haipengluo/.conda/envs/sdxl/bin/uvicorn --host=phlrr3006.guest.corp.microsoft.com --port 7999 --workers 1 --backlog 1 --limit-concurrency 4 main_v3:app
+v-haipe+ 13861 13722  0 Sep24 pts/6    00:00:05 /home/v-haipengluo/.conda/envs/sdxl/bin/python /home/v-haipengluo/.conda/envs/sdxl/bin/uvicorn --host=phlrr3006.guest.corp.microsoft.com --port 7999 --workers 1 --backlog 1 --limit-concurrency 4 main_v3:app
+v-haipe+ 13863 13722  0 Sep24 pts/6    00:00:05 /home/v-haipengluo/.conda/envs/sdxl/bin/python /home/v-haipengluo/.conda/envs/sdxl/bin/uvicorn --host=phlrr3006.guest.corp.microsoft.com --port 7999 --workers 1 --backlog 1 --limit-concurrency 4 main_v3:app
+v-haipe+ 13865 13722  0 Sep24 pts/6    00:00:05 /home/v-haipengluo/.conda/envs/sdxl/bin/python /home/v-haipengluo/.conda/envs/sdxl/bin/uvicorn --host=phlrr3006.guest.corp.microsoft.com --port 7999 --workers 1 --backlog 1 --limit-concurrency 4 main_v3:app
+v-haipe+ 13867 13722  0 Sep24 pts/6    00:00:05 /home/v-haipengluo/.conda/envs/sdxl/bin/python /home/v-haipengluo/.conda/envs/sdxl/bin/uvicorn --host=phlrr3006.guest.corp.microsoft.com --port 7999 --workers 1 --backlog 1 --limit-concurrency 4 main_v3:app
+v-haipe+ 13869 13722  0 Sep24 pts/6    00:00:05 /home/v-haipengluo/.conda/envs/sdxl/bin/python /home/v-haipengluo/.conda/envs/sdxl/bin/uvicorn --host=phlrr3006.guest.corp.microsoft.com --port 7999 --workers 1 --backlog 1 --limit-concurrency 4 main_v3:app
+v-haipe+ 13871 13722  0 Sep24 pts/6    00:00:05 /home/v-haipengluo/.conda/envs/sdxl/bin/python /home/v-haipengluo/.conda/envs/sdxl/bin/uvicorn --host=phlrr3006.guest.corp.microsoft.com --port 7999 --workers 1 --backlog 1 --limit-concurrency 4 main_v3:app
+v-haipe+ 13873 13722  0 Sep24 pts/6    00:00:05 /home/v-haipengluo/.conda/envs/sdxl/bin/python /home/v-haipengluo/.conda/envs/sdxl/bin/uvicorn --host=phlrr3006.guest.corp.microsoft.com --port 7999 --workers 1 --backlog 1 --limit-concurrency 4 main_v3:app
+v-haipe+ 13875 13722  0 Sep24 pts/6    00:00:05 /home/v-haipengluo/.conda/envs/sdxl/bin/python /home/v-haipengluo/.conda/envs/sdxl/bin/uvicorn --host=phlrr3006.guest.corp.microsoft.com --port 7999 --workers 1 --backlog 1 --limit-concurrency 4 main_v3:app
+v-haipe+ 13877 13722  0 Sep24 pts/6    00:00:05 /home/v-haipengluo/.conda/envs/sdxl/bin/python /home/v-haipengluo/.conda/envs/sdxl/bin/uvicorn --host=phlrr3006.guest.corp.microsoft.com --port 7999 --workers 1 --backlog 1 --limit-concurrency 4 main_v3:app
+v-haipe+ 13879 13722  0 Sep24 pts/6    00:00:05 /home/v-haipengluo/.conda/envs/sdxl/bin/python /home/v-haipengluo/.conda/envs/sdxl/bin/uvicorn --host=phlrr3006.guest.corp.microsoft.com --port 7999 --workers 1 --backlog 1 --limit-concurrency 4 main_v3:app
+v-haipe+ 13881 13722  0 Sep24 pts/6    00:00:05 /home/v-haipengluo/.conda/envs/sdxl/bin/python /home/v-haipengluo/.conda/envs/sdxl/bin/uvicorn --host=phlrr3006.guest.corp.microsoft.com --port 7999 --workers 1 --backlog 1 --limit-concurrency 4 main_v3:app
+v-haipe+ 13883 13722  0 Sep24 pts/6    00:00:05 /home/v-haipengluo/.conda/envs/sdxl/bin/python /home/v-haipengluo/.conda/envs/sdxl/bin/uvicorn --host=phlrr3006.guest.corp.microsoft.com --port 7999 --workers 1 --backlog 1 --limit-concurrency 4 main_v3:app
+v-haipe+ 13885 13722  0 Sep24 pts/6    00:00:05 /home/v-haipengluo/.conda/envs/sdxl/bin/python /home/v-haipengluo/.conda/envs/sdxl/bin/uvicorn --host=phlrr3006.guest.corp.microsoft.com --port 7999 --workers 1 --backlog 1 --limit-concurrency 4 main_v3:app
+v-haipe+ 13887 13722  0 Sep24 pts/6    00:00:05 /home/v-haipengluo/.conda/envs/sdxl/bin/python /home/v-haipengluo/.conda/envs/sdxl/bin/uvicorn --host=phlrr3006.guest.corp.microsoft.com --port 7999 --workers 1 --backlog 1 --limit-concurrency 4 main_v3:app
+v-haipe+ 18319 15852  0 05:34 pts/1    00:00:03 /home/v-haipengluo/.conda/envs/llamax/bin/python /home/v-haipengluo/.conda/envs/llamax/bin/deepspeed --master_port 29500 --hostfile=hostfile --include=localhost:1,3,4,5,6,7 src/train.py --model_name_or_path /workspaceblobstore/qins/test/20220316/haipeng/output_weights/llamax_13b_stackexchange_MATH_12w_sample_5w_score0.5_trainset_2e-5/checkpoint-992 --data_path /workspaceblobstore/qins/test/20220316/haipeng/data/Math_datasets/MATH_the_answer_is_format/hendrycks_math_7500_ori_gpt4_ori_15k.json --output_dir /workspaceblobstore/qins/test/20220316/haipeng/output_weights/llamax_13b_continue_train_stackMATH5w_checkpoint992_hendrycks_math_7500_ori_gpt4_ori_15k --num_train_epochs 3 --model_max_length 1150 --per_device_train_batch_size 17 --per_device_eval_batch_size 1 --gradient_accumulation_steps 1 --evaluation_strategy no --save_strategy steps --save_steps 36 --save_total_limit 200 --learning_rate 2e-5 --warmup_steps 10 --logging_steps 2 --lr_scheduler_type cosine --report_to tensorboard --gradient_checkpointing True --deepspeed src/configs/deepspeed_config.json --fp16 True
+v-haipe+ 18333 18319  0 05:34 pts/1    00:00:03 /home/v-haipengluo/.conda/envs/llamax/bin/python -u -m deepspeed.launcher.launch --world_info=eyJsb2NhbGhvc3QiOiBbMSwgMywgNCwgNSwgNiwgN119 --master_addr=127.0.0.1 --master_port=29500 --enable_each_rank_log=None src/train.py --model_name_or_path /workspaceblobstore/qins/test/20220316/haipeng/output_weights/llamax_13b_stackexchange_MATH_12w_sample_5w_score0.5_trainset_2e-5/checkpoint-992 --data_path /workspaceblobstore/qins/test/20220316/haipeng/data/Math_datasets/MATH_the_answer_is_format/hendrycks_math_7500_ori_gpt4_ori_15k.json --output_dir /workspaceblobstore/qins/test/20220316/haipeng/output_weights/llamax_13b_continue_train_stackMATH5w_checkpoint992_hendrycks_math_7500_ori_gpt4_ori_15k --num_train_epochs 3 --model_max_length 1150 --per_device_train_batch_size 17 --per_device_eval_batch_size 1 --gradient_accumulation_steps 1 --evaluation_strategy no --save_strategy steps --save_steps 36 --save_total_limit 200 --learning_rate 2e-5 --warmup_steps 10 --logging_steps 2 --lr_scheduler_type cosine --report_to tensorboard --gradient_checkpointing True --deepspeed src/configs/deepspeed_config.json --fp16 True
+v-haipe+ 18346 18333 99 05:34 pts/1    03:20:42 /home/v-haipengluo/.conda/envs/llamax/bin/python -u src/train.py --local_rank=0 --model_name_or_path /workspaceblobstore/qins/test/20220316/haipeng/output_weights/llamax_13b_stackexchange_MATH_12w_sample_5w_score0.5_trainset_2e-5/checkpoint-992 --data_path /workspaceblobstore/qins/test/20220316/haipeng/data/Math_datasets/MATH_the_answer_is_format/hendrycks_math_7500_ori_gpt4_ori_15k.json --output_dir /workspaceblobstore/qins/test/20220316/haipeng/output_weights/llamax_13b_continue_train_stackMATH5w_checkpoint992_hendrycks_math_7500_ori_gpt4_ori_15k --num_train_epochs 3 --model_max_length 1150 --per_device_train_batch_size 17 --per_device_eval_batch_size 1 --gradient_accumulation_steps 1 --evaluation_strategy no --save_strategy steps --save_steps 36 --save_total_limit 200 --learning_rate 2e-5 --warmup_steps 10 --logging_steps 2 --lr_scheduler_type cosine --report_to tensorboard --gradient_checkpointing True --deepspeed src/configs/deepspeed_config.json --fp16 True
+v-haipe+ 18347 18333 99 05:34 pts/1    03:40:59 /home/v-haipengluo/.conda/envs/llamax/bin/python -u src/train.py --local_rank=1 --model_name_or_path /workspaceblobstore/qins/test/20220316/haipeng/output_weights/llamax_13b_stackexchange_MATH_12w_sample_5w_score0.5_trainset_2e-5/checkpoint-992 --data_path /workspaceblobstore/qins/test/20220316/haipeng/data/Math_datasets/MATH_the_answer_is_format/hendrycks_math_7500_ori_gpt4_ori_15k.json --output_dir /workspaceblobstore/qins/test/20220316/haipeng/output_weights/llamax_13b_continue_train_stackMATH5w_checkpoint992_hendrycks_math_7500_ori_gpt4_ori_15k --num_train_epochs 3 --model_max_length 1150 --per_device_train_batch_size 17 --per_device_eval_batch_size 1 --gradient_accumulation_steps 1 --evaluation_strategy no --save_strategy steps --save_steps 36 --save_total_limit 200 --learning_rate 2e-5 --warmup_steps 10 --logging_steps 2 --lr_scheduler_type cosine --report_to tensorboard --gradient_checkpointing True --deepspeed src/configs/deepspeed_config.json --fp16 True
+v-haipe+ 18348 18333 99 05:34 pts/1    03:44:08 /home/v-haipengluo/.conda/envs/llamax/bin/python -u src/train.py --local_rank=2 --model_name_or_path /workspaceblobstore/qins/test/20220316/haipeng/output_weights/llamax_13b_stackexchange_MATH_12w_sample_5w_score0.5_trainset_2e-5/checkpoint-992 --data_path /workspaceblobstore/qins/test/20220316/haipeng/data/Math_datasets/MATH_the_answer_is_format/hendrycks_math_7500_ori_gpt4_ori_15k.json --output_dir /workspaceblobstore/qins/test/20220316/haipeng/output_weights/llamax_13b_continue_train_stackMATH5w_checkpoint992_hendrycks_math_7500_ori_gpt4_ori_15k --num_train_epochs 3 --model_max_length 1150 --per_device_train_batch_size 17 --per_device_eval_batch_size 1 --gradient_accumulation_steps 1 --evaluation_strategy no --save_strategy steps --save_steps 36 --save_total_limit 200 --learning_rate 2e-5 --warmup_steps 10 --logging_steps 2 --lr_scheduler_type cosine --report_to tensorboard --gradient_checkpointing True --deepspeed src/configs/deepspeed_config.json --fp16 True
+v-haipe+ 18349 18333 99 05:34 pts/1    03:32:51 /home/v-haipengluo/.conda/envs/llamax/bin/python -u src/train.py --local_rank=3 --model_name_or_path /workspaceblobstore/qins/test/20220316/haipeng/output_weights/llamax_13b_stackexchange_MATH_12w_sample_5w_score0.5_trainset_2e-5/checkpoint-992 --data_path /workspaceblobstore/qins/test/20220316/haipeng/data/Math_datasets/MATH_the_answer_is_format/hendrycks_math_7500_ori_gpt4_ori_15k.json --output_dir /workspaceblobstore/qins/test/20220316/haipeng/output_weights/llamax_13b_continue_train_stackMATH5w_checkpoint992_hendrycks_math_7500_ori_gpt4_ori_15k --num_train_epochs 3 --model_max_length 1150 --per_device_train_batch_size 17 --per_device_eval_batch_size 1 --gradient_accumulation_steps 1 --evaluation_strategy no --save_strategy steps --save_steps 36 --save_total_limit 200 --learning_rate 2e-5 --warmup_steps 10 --logging_steps 2 --lr_scheduler_type cosine --report_to tensorboard --gradient_checkpointing True --deepspeed src/configs/deepspeed_config.json --fp16 True
+v-haipe+ 18350 18333 99 05:34 pts/1    03:41:16 /home/v-haipengluo/.conda/envs/llamax/bin/python -u src/train.py --local_rank=4 --model_name_or_path /workspaceblobstore/qins/test/20220316/haipeng/output_weights/llamax_13b_stackexchange_MATH_12w_sample_5w_score0.5_trainset_2e-5/checkpoint-992 --data_path /workspaceblobstore/qins/test/20220316/haipeng/data/Math_datasets/MATH_the_answer_is_format/hendrycks_math_7500_ori_gpt4_ori_15k.json --output_dir /workspaceblobstore/qins/test/20220316/haipeng/output_weights/llamax_13b_continue_train_stackMATH5w_checkpoint992_hendrycks_math_7500_ori_gpt4_ori_15k --num_train_epochs 3 --model_max_length 1150 --per_device_train_batch_size 17 --per_device_eval_batch_size 1 --gradient_accumulation_steps 1 --evaluation_strategy no --save_strategy steps --save_steps 36 --save_total_limit 200 --learning_rate 2e-5 --warmup_steps 10 --logging_steps 2 --lr_scheduler_type cosine --report_to tensorboard --gradient_checkpointing True --deepspeed src/configs/deepspeed_config.json --fp16 True
+v-haipe+ 18351 18333 99 05:34 pts/1    03:42:27 /home/v-haipengluo/.conda/envs/llamax/bin/python -u src/train.py --local_rank=5 --model_name_or_path /workspaceblobstore/qins/test/20220316/haipeng/output_weights/llamax_13b_stackexchange_MATH_12w_sample_5w_score0.5_trainset_2e-5/checkpoint-992 --data_path /workspaceblobstore/qins/test/20220316/haipeng/data/Math_datasets/MATH_the_answer_is_format/hendrycks_math_7500_ori_gpt4_ori_15k.json --output_dir /workspaceblobstore/qins/test/20220316/haipeng/output_weights/llamax_13b_continue_train_stackMATH5w_checkpoint992_hendrycks_math_7500_ori_gpt4_ori_15k --num_train_epochs 3 --model_max_length 1150 --per_device_train_batch_size 17 --per_device_eval_batch_size 1 --gradient_accumulation_steps 1 --evaluation_strategy no --save_strategy steps --save_steps 36 --save_total_limit 200 --learning_rate 2e-5 --warmup_steps 10 --logging_steps 2 --lr_scheduler_type cosine --report_to tensorboard --gradient_checkpointing True --deepspeed src/configs/deepspeed_config.json --fp16 True
+v-haipe+ 24334 23818  0 Sep23 pts/7    00:00:25 python -m http.server

img/stable-diffusion-server/main.py

@@ -0,0 +1,528 @@
+import gc
+import math
+import multiprocessing
+import os
+import traceback
+from datetime import datetime
+from io import BytesIO
+from itertools import permutations
+from multiprocessing.pool import Pool
+from pathlib import Path
+from urllib.parse import quote_plus
+
+import numpy as np
+import nltk
+import torch
+
+from PIL.Image import Image
+from diffusers import DiffusionPipeline, StableDiffusionXLInpaintPipeline
+from diffusers.utils import load_image
+from fastapi import FastAPI
+from fastapi.middleware.gzip import GZipMiddleware
+from loguru import logger
+from starlette.middleware.cors import CORSMiddleware
+from starlette.responses import FileResponse
+from starlette.responses import JSONResponse
+
+from env import BUCKET_PATH, BUCKET_NAME
+# from stable_diffusion_server.bucket_api import check_if_blob_exists, upload_to_bucket
+torch._dynamo.config.suppress_errors = True
+
+import string
+import random
+
+def generate_save_path():
+    # initializing size of string
+    N = 7
+
+    # using random.choices()
+    # generating random strings
+    res = ''.join(random.choices(string.ascii_uppercase +
+                                 string.digits, k=N))
+    return res
+
+pipe = DiffusionPipeline.from_pretrained(
+    "models/stable-diffusion-xl-base-1.0",
+    torch_dtype=torch.bfloat16,
+    use_safetensors=True,
+    variant="fp16",
+    # safety_checker=None,
+)  # todo try torch_dtype=bfloat16
+pipe.watermark = None
+
+pipe.to("cuda")
+
+refiner = DiffusionPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-refiner-1.0",
+    text_encoder_2=pipe.text_encoder_2,
+    vae=pipe.vae,
+    torch_dtype=torch.bfloat16, # safer to use bfloat?
+    use_safetensors=True,
+    variant="fp16", #remember not to download the big model
+)
+refiner.watermark = None
+refiner.to("cuda")
+
+# {'scheduler', 'text_encoder', 'text_encoder_2', 'tokenizer', 'tokenizer_2', 'unet', 'vae'} can be passed in from existing model
+inpaintpipe = StableDiffusionXLInpaintPipeline.from_pretrained(
+    "models/stable-diffusion-xl-base-1.0", torch_dtype=torch.bfloat16, variant="fp16", use_safetensors=True,
+    scheduler=pipe.scheduler,
+    text_encoder=pipe.text_encoder,
+    text_encoder_2=pipe.text_encoder_2,
+    tokenizer=pipe.tokenizer,
+    tokenizer_2=pipe.tokenizer_2,
+    unet=pipe.unet,
+    vae=pipe.vae,
+    # load_connected_pipeline=
+)
+# # switch out to save gpu mem
+# del inpaintpipe.vae
+# del inpaintpipe.text_encoder_2
+# del inpaintpipe.text_encoder
+# del inpaintpipe.scheduler
+# del inpaintpipe.tokenizer
+# del inpaintpipe.tokenizer_2
+# del inpaintpipe.unet
+# inpaintpipe.vae = pipe.vae
+# inpaintpipe.text_encoder_2 = pipe.text_encoder_2
+# inpaintpipe.text_encoder = pipe.text_encoder
+# inpaintpipe.scheduler = pipe.scheduler
+# inpaintpipe.tokenizer = pipe.tokenizer
+# inpaintpipe.tokenizer_2 = pipe.tokenizer_2
+# inpaintpipe.unet = pipe.unet
+# todo this should work
+# inpaintpipe = StableDiffusionXLInpaintPipeline( # construct an inpainter using the existing model
+#     vae=pipe.vae,
+#     text_encoder_2=pipe.text_encoder_2,
+#     text_encoder=pipe.text_encoder,
+#     unet=pipe.unet,
+#     scheduler=pipe.scheduler,
+#     tokenizer=pipe.tokenizer,
+#     tokenizer_2=pipe.tokenizer_2,
+#     requires_aesthetics_score=False,
+# )
+inpaintpipe.to("cuda")
+inpaintpipe.watermark = None
+# inpaintpipe.register_to_config(requires_aesthetics_score=False)
+
+inpaint_refiner = StableDiffusionXLInpaintPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-refiner-1.0",
+    text_encoder_2=inpaintpipe.text_encoder_2,
+    vae=inpaintpipe.vae,
+    torch_dtype=torch.bfloat16,
+    use_safetensors=True,
+    variant="fp16",
+
+    tokenizer_2=refiner.tokenizer_2,
+    tokenizer=refiner.tokenizer,
+    scheduler=refiner.scheduler,
+    text_encoder=refiner.text_encoder,
+    unet=refiner.unet,
+)
+# del inpaint_refiner.vae
+# del inpaint_refiner.text_encoder_2
+# del inpaint_refiner.text_encoder
+# del inpaint_refiner.scheduler
+# del inpaint_refiner.tokenizer
+# del inpaint_refiner.tokenizer_2
+# del inpaint_refiner.unet
+# inpaint_refiner.vae = inpaintpipe.vae
+# inpaint_refiner.text_encoder_2 = inpaintpipe.text_encoder_2
+#
+# inpaint_refiner.text_encoder = refiner.text_encoder
+# inpaint_refiner.scheduler = refiner.scheduler
+# inpaint_refiner.tokenizer = refiner.tokenizer
+# inpaint_refiner.tokenizer_2 = refiner.tokenizer_2
+# inpaint_refiner.unet = refiner.unet
+
+# inpaint_refiner = StableDiffusionXLInpaintPipeline(
+#     text_encoder_2=inpaintpipe.text_encoder_2,
+#     vae=inpaintpipe.vae,
+#     # the rest from the existing refiner
+#     tokenizer_2=refiner.tokenizer_2,
+#     tokenizer=refiner.tokenizer,
+#     scheduler=refiner.scheduler,
+#     text_encoder=refiner.text_encoder,
+#     unet=refiner.unet,
+#     requires_aesthetics_score=False,
+# )
+inpaint_refiner.to("cuda")
+inpaint_refiner.watermark = None
+# inpaint_refiner.register_to_config(requires_aesthetics_score=False)
+
+n_steps = 40
+high_noise_frac = 0.8
+
+# if using torch < 2.0
+# pipe.enable_xformers_memory_efficient_attention()
+
+
+# pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
+# this can cause errors on some inputs so consider disabling it
+pipe.unet = torch.compile(pipe.unet)
+refiner.unet = torch.compile(refiner.unet)#, mode="reduce-overhead", fullgraph=True)
+# compile the inpainters - todo reuse the other unets? swap out the models for others/del them so they share models and can be swapped efficiently
+inpaintpipe.unet = pipe.unet
+inpaint_refiner.unet = refiner.unet
+# inpaintpipe.unet = torch.compile(inpaintpipe.unet)
+# inpaint_refiner.unet = torch.compile(inpaint_refiner.unet)
+from pydantic import BaseModel
+
+app = FastAPI(
+    openapi_url="/static/openapi.json",
+    docs_url="/swagger-docs",
+    redoc_url="/redoc",
+    title="Generate Images Netwrck API",
+    description="Character Chat API",
+    # root_path="https://api.text-generator.io",
+    version="1",
+)
+app.add_middleware(GZipMiddleware, minimum_size=1000)
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+stopwords = nltk.corpus.stopwords.words("english")
+
+class Img(BaseModel):
+    system_prompt: str
+    ASSISTANT: str
+
+# img_url = "http://phlrr2019.guest.corp.microsoft.com:8000/img1_sdv2.1.png"
+img_url = "http://phlrr3058.guest.corp.microsoft.com:8000/"#/img1_sdv2.1.png"
+
+@app.post("/image_url")
+def image_url(img: Img):
+    system_prompt = img.system_prompt
+    prompt = img.ASSISTANT
+    # if Path(save_path).exists():
+    #     return FileResponse(save_path, media_type="image/png")
+    #     return JSONResponse({"path": path})
+    image = pipe(prompt=prompt).images[0]
+    # if not save_path:
+    save_path = generate_save_path()
+    save_path = f"images/{save_path}.png"
+    image.save(save_path)
+    # save_path = '/'.join(path_components) + quote_plus(final_name)
+    path = f"{img_url}/{save_path}"
+    return JSONResponse({"path": path})
+
+
+@app.get("/make_image")
+# @app.post("/make_image")
+def make_image(prompt: str, save_path: str = ""):
+    if Path(save_path).exists():
+        return FileResponse(save_path, media_type="image/png")
+    image = pipe(prompt=prompt).images[0]
+    if not save_path:
+        save_path = f"images/{prompt}.png"
+    image.save(save_path)
+    return FileResponse(save_path, media_type="image/png")
+
+
+@app.get("/create_and_upload_image")
+def create_and_upload_image(prompt: str, width: int=1024, height:int=1024, save_path: str = ""):
+    path_components = save_path.split("/")[0:-1]
+    final_name = save_path.split("/")[-1]
+    if not path_components:
+        path_components = []
+    save_path = '/'.join(path_components) + quote_plus(final_name)
+    path = get_image_or_create_upload_to_cloud_storage(prompt, width, height, save_path)
+    return JSONResponse({"path": path})
+
+@app.get("/inpaint_and_upload_image")
+def inpaint_and_upload_image(prompt: str, image_url:str, mask_url:str, save_path: str = ""):
+    path_components = save_path.split("/")[0:-1]
+    final_name = save_path.split("/")[-1]
+    if not path_components:
+        path_components = []
+    save_path = '/'.join(path_components) + quote_plus(final_name)
+    path = get_image_or_inpaint_upload_to_cloud_storage(prompt, image_url, mask_url, save_path)
+    return JSONResponse({"path": path})
+
+
+def get_image_or_create_upload_to_cloud_storage(prompt:str,width:int, height:int, save_path:str):
+    prompt = shorten_too_long_text(prompt)
+    save_path = shorten_too_long_text(save_path)
+    # check exists - todo cache this
+    if check_if_blob_exists(save_path):
+        return f"https://{BUCKET_NAME}/{BUCKET_PATH}/{save_path}"
+    bio = create_image_from_prompt(prompt, width, height)
+    if bio is None:
+        return None # error thrown in pool
+    link = upload_to_bucket(save_path, bio, is_bytesio=True)
+    return link
+def get_image_or_inpaint_upload_to_cloud_storage(prompt:str, image_url:str, mask_url:str, save_path:str):
+    prompt = shorten_too_long_text(prompt)
+    save_path = shorten_too_long_text(save_path)
+    # check exists - todo cache this
+    if check_if_blob_exists(save_path):
+        return f"https://{BUCKET_NAME}/{BUCKET_PATH}/{save_path}"
+    bio = inpaint_image_from_prompt(prompt, image_url, mask_url)
+    if bio is None:
+        return None # error thrown in pool
+    link = upload_to_bucket(save_path, bio, is_bytesio=True)
+    return link
+
+# multiprocessing.set_start_method('spawn', True)
+# processes_pool = Pool(1) # cant do too much at once or OOM errors happen
+# def create_image_from_prompt_sync(prompt):
+#     """have to call this sync to avoid OOM errors"""
+#     return processes_pool.apply_async(create_image_from_prompt, args=(prompt,), ).wait()
+
+def create_image_from_prompt(prompt, width, height):
+    # round width and height down to multiple of 64
+    block_width = width - (width % 64)
+    block_height = height - (height % 64)
+    prompt = shorten_too_long_text(prompt)
+    # image = pipe(prompt=prompt).images[0]
+    try:
+        image = pipe(prompt=prompt,
+                     width=block_width,
+                     height=block_height,
+                     # denoising_end=high_noise_frac,
+                     # output_type='latent',
+                     # height=512,
+                     # width=512,
+                     num_inference_steps=50).images[0]  # normally uses 50 steps
+    except Exception as e:
+        # try rm stopwords + half the prompt
+        # todo try prompt permutations
+        logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+        prompt = ' '.join((word for word in prompt if word not in stopwords))
+        prompts = prompt.split()
+
+        prompt = ' '.join(prompts[:len(prompts) // 2])
+        logger.info(f"shortened prompt to: {len(prompt)}")
+        image = None
+        if prompt:
+            try:
+                image = pipe(prompt=prompt,
+                             width=block_width,
+                             height=block_height,
+                             # denoising_end=high_noise_frac,
+                             # output_type='latent',
+                             # height=512,
+                             # width=512,
+                             num_inference_steps=50).images[0]  # normally uses 50 steps
+            except Exception as e:
+                # logger.info("trying to permute prompt")
+                # # try two swaps of the prompt/permutations
+                # prompt = prompt.split()
+                # prompt = ' '.join(permutations(prompt, 2).__next__())
+                logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+                prompt = ' '.join((word for word in prompt if word not in stopwords))
+                prompts = prompt.split()
+
+                prompt = ' '.join(prompts[:len(prompts) // 2])
+                logger.info(f"shortened prompt to: {len(prompt)}")
+
+                try:
+                    image = pipe(prompt=prompt,
+                                 width=block_width,
+                                 height=block_height,
+                                 # denoising_end=high_noise_frac,
+                                 # output_type='latent', # dont need latent yet - we refine the image at full res
+                                 # height=512,
+                                 # width=512,
+                                 num_inference_steps=50).images[0]  # normally uses 50 steps
+                except Exception as e:
+                    # just error out
+                    traceback.print_exc()
+                    raise e
+                    # logger.info("restarting server to fix cuda issues (device side asserts)")
+                    # todo fix device side asserts instead of restart to fix
+                    # todo only restart the correct gunicorn
+                    # this could be really annoying if your running other gunicorns on your machine which also get restarted
+                    # os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+                    # os.system("kill -1 `pgrep gunicorn`")
+    # todo refine
+    # if image != None:
+    #     image = refiner(
+    #         prompt=prompt,
+    #         # width=block_width,
+    #         # height=block_height,
+    #         num_inference_steps=n_steps,
+    #         # denoising_start=high_noise_frac,
+    #         image=image,
+    #     ).images[0]
+    if width != block_width or height != block_height:
+        # resize to original size width/height
+        # find aspect ratio to scale up to that covers the original img input width/height
+        scale_up_ratio = max(width / block_width, height / block_height)
+        image = image.resize((math.ceil(block_width * scale_up_ratio), math.ceil(height * scale_up_ratio)))
+        # crop image to original size
+        image = image.crop((0, 0, width, height))
+    # try:
+    #     # gc.collect()
+    #     torch.cuda.empty_cache()
+    # except Exception as e:
+    #     traceback.print_exc()
+    #     logger.info("restarting server to fix cuda issues (device side asserts)")
+    #     # todo fix device side asserts instead of restart to fix
+    #     # todo only restart the correct gunicorn
+    #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+    #     os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+    #     os.system("kill -1 `pgrep gunicorn`")
+    # save as bytesio
+    bs = BytesIO()
+
+    bright_count = np.sum(np.array(image) > 0)
+    if bright_count == 0:
+        # we have a black image, this is an error likely we need a restart
+        logger.info("restarting server to fix cuda issues (device side asserts)")
+        #     # todo fix device side asserts instead of restart to fix
+        #     # todo only restart the correct gunicorn
+        #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+        os.system("kill -1 `pgrep gunicorn`")
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep uvicorn`")
+        os.system("kill -1 `pgrep uvicorn`")
+
+        return None
+    image.save(bs, quality=85, optimize=True, format="webp")
+    bio = bs.getvalue()
+    # touch progress.txt file - if we dont do this we get restarted by supervisor/other processes for reliability
+    with open("progress.txt", "w") as f:
+        current_time = datetime.now().strftime("%H:%M:%S")
+        f.write(f"{current_time}")
+    return bio
+
+def inpaint_image_from_prompt(prompt, image_url: str, mask_url: str):
+    prompt = shorten_too_long_text(prompt)
+    # image = pipe(prompt=prompt).images[0]
+
+    init_image = load_image(image_url).convert("RGB")
+    mask_image = load_image(mask_url).convert("RGB") # why rgb for a 1 channel mask?
+    num_inference_steps = 75
+    high_noise_frac = 0.7
+
+    try:
+        image = inpaintpipe(
+            prompt=prompt,
+            image=init_image,
+            mask_image=mask_image,
+            num_inference_steps=num_inference_steps,
+            denoising_start=high_noise_frac,
+            output_type="latent",
+        ).images[0]  # normally uses 50 steps
+    except Exception as e:
+        # try rm stopwords + half the prompt
+        # todo try prompt permutations
+        logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+        prompt = ' '.join((word for word in prompt if word not in stopwords))
+        prompts = prompt.split()
+
+        prompt = ' '.join(prompts[:len(prompts) // 2])
+        logger.info(f"shortened prompt to: {len(prompt)}")
+        image = None
+        if prompt:
+            try:
+                image = pipe(
+                    prompt=prompt,
+                    image=init_image,
+                    mask_image=mask_image,
+                    num_inference_steps=num_inference_steps,
+                    denoising_start=high_noise_frac,
+                    output_type="latent",
+                ).images[0]  # normally uses 50 steps
+            except Exception as e:
+                # logger.info("trying to permute prompt")
+                # # try two swaps of the prompt/permutations
+                # prompt = prompt.split()
+                # prompt = ' '.join(permutations(prompt, 2).__next__())
+                logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+                prompt = ' '.join((word for word in prompt if word not in stopwords))
+                prompts = prompt.split()
+
+                prompt = ' '.join(prompts[:len(prompts) // 2])
+                logger.info(f"shortened prompt to: {len(prompt)}")
+
+                try:
+                    image = inpaintpipe(
+                        prompt=prompt,
+                        image=init_image,
+                        mask_image=mask_image,
+                        num_inference_steps=num_inference_steps,
+                        denoising_start=high_noise_frac,
+                        output_type="latent",
+                    ).images[0]  # normally uses 50 steps
+                except Exception as e:
+                    # just error out
+                    traceback.print_exc()
+                    raise e
+                    # logger.info("restarting server to fix cuda issues (device side asserts)")
+                    # todo fix device side asserts instead of restart to fix
+                    # todo only restart the correct gunicorn
+                    # this could be really annoying if your running other gunicorns on your machine which also get restarted
+                    # os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+                    # os.system("kill -1 `pgrep gunicorn`")
+    if image != None:
+        image = inpaint_refiner(
+            prompt=prompt,
+            image=image,
+            mask_image=mask_image,
+            num_inference_steps=num_inference_steps,
+            denoising_start=high_noise_frac,
+
+        ).images[0]
+    # try:
+    #     # gc.collect()
+    #     torch.cuda.empty_cache()
+    # except Exception as e:
+    #     traceback.print_exc()
+    #     logger.info("restarting server to fix cuda issues (device side asserts)")
+    #     # todo fix device side asserts instead of restart to fix
+    #     # todo only restart the correct gunicorn
+    #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+    #     os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+    #     os.system("kill -1 `pgrep gunicorn`")
+    # save as bytesio
+    bs = BytesIO()
+
+    bright_count = np.sum(np.array(image) > 0)
+    if bright_count == 0:
+        # we have a black image, this is an error likely we need a restart
+        logger.info("restarting server to fix cuda issues (device side asserts)")
+        #     # todo fix device side asserts instead of restart to fix
+        #     # todo only restart the correct gunicorn
+        #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+        os.system("kill -1 `pgrep gunicorn`")
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep uvicorn`")
+        os.system("kill -1 `pgrep uvicorn`")
+
+        return None
+    image.save(bs, quality=85, optimize=True, format="webp")
+    bio = bs.getvalue()
+    # touch progress.txt file - if we dont do this we get restarted by supervisor/other processes for reliability
+    with open("progress.txt", "w") as f:
+        current_time = datetime.now().strftime("%H:%M:%S")
+        f.write(f"{current_time}")
+    return bio
+
+
+
+def shorten_too_long_text(prompt):
+    if len(prompt) > 200:
+        # remove stopwords
+        prompt = prompt.split() # todo also split hyphens
+        prompt = ' '.join((word for word in prompt if word not in stopwords))
+        if len(prompt) > 200:
+            prompt = prompt[:200]
+    return prompt
+
+# image = pipe(prompt=prompt).images[0]
+#
+# image.save("test.png")
+# # save all images
+# for i, image in enumerate(images):
+#     image.save(f"{i}.png")

img/stable-diffusion-server/main_1024.py

@@ -0,0 +1,549 @@
+import gc
+import math
+import multiprocessing
+import os
+import traceback
+from datetime import datetime
+from io import BytesIO
+from itertools import permutations
+from multiprocessing.pool import Pool
+from pathlib import Path
+from urllib.parse import quote_plus
+
+import numpy as np
+import nltk
+import torch
+
+from PIL.Image import Image
+from diffusers import DiffusionPipeline, StableDiffusionXLInpaintPipeline
+from diffusers.utils import load_image
+from fastapi import FastAPI
+from fastapi.middleware.gzip import GZipMiddleware
+from loguru import logger
+from starlette.middleware.cors import CORSMiddleware
+from starlette.responses import FileResponse
+from starlette.responses import JSONResponse
+
+from env import BUCKET_PATH, BUCKET_NAME
+# from stable_diffusion_server.bucket_api import check_if_blob_exists, upload_to_bucket
+torch._dynamo.config.suppress_errors = True
+
+import string
+import random
+
+def generate_save_path():
+    # initializing size of string
+    N = 7
+
+    # using random.choices()
+    # generating random strings
+    res = ''.join(random.choices(string.ascii_uppercase +
+                                 string.digits, k=N))
+    return res
+
+# pipe = DiffusionPipeline.from_pretrained(
+#     "models/stable-diffusion-xl-base-1.0",
+#     torch_dtype=torch.bfloat16,
+#     use_safetensors=True,
+#     variant="fp16",
+#     # safety_checker=None,
+# )  # todo try torch_dtype=bfloat16
+
+model_dir = os.getenv("SDXL_MODEL_DIR")
+
+if model_dir:
+    # Use local model
+    model_key_base = os.path.join(model_dir, "stable-diffusion-xl-base-1.0")
+    model_key_refiner = os.path.join(model_dir, "stable-diffusion-xl-refiner-1.0")
+else:
+    model_key_base = "stabilityai/stable-diffusion-xl-base-1.0"
+    model_key_refiner = "stabilityai/stable-diffusion-xl-refiner-1.0"
+
+pipe = DiffusionPipeline.from_pretrained(model_key_base, torch_dtype=torch.float16, use_safetensors=True, variant="fp16")
+
+pipe.watermark = None
+
+pipe.to("cuda")
+
+refiner = DiffusionPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-refiner-1.0",
+    text_encoder_2=pipe.text_encoder_2,
+    vae=pipe.vae,
+    torch_dtype=torch.bfloat16, # safer to use bfloat?
+    use_safetensors=True,
+    variant="fp16", #remember not to download the big model
+)
+refiner.watermark = None
+refiner.to("cuda")
+
+# {'scheduler', 'text_encoder', 'text_encoder_2', 'tokenizer', 'tokenizer_2', 'unet', 'vae'} can be passed in from existing model
+inpaintpipe = StableDiffusionXLInpaintPipeline.from_pretrained(
+    "models/stable-diffusion-xl-base-1.0", torch_dtype=torch.bfloat16, variant="fp16", use_safetensors=True,
+    scheduler=pipe.scheduler,
+    text_encoder=pipe.text_encoder,
+    text_encoder_2=pipe.text_encoder_2,
+    tokenizer=pipe.tokenizer,
+    tokenizer_2=pipe.tokenizer_2,
+    unet=pipe.unet,
+    vae=pipe.vae,
+    # load_connected_pipeline=
+)
+# # switch out to save gpu mem
+# del inpaintpipe.vae
+# del inpaintpipe.text_encoder_2
+# del inpaintpipe.text_encoder
+# del inpaintpipe.scheduler
+# del inpaintpipe.tokenizer
+# del inpaintpipe.tokenizer_2
+# del inpaintpipe.unet
+# inpaintpipe.vae = pipe.vae
+# inpaintpipe.text_encoder_2 = pipe.text_encoder_2
+# inpaintpipe.text_encoder = pipe.text_encoder
+# inpaintpipe.scheduler = pipe.scheduler
+# inpaintpipe.tokenizer = pipe.tokenizer
+# inpaintpipe.tokenizer_2 = pipe.tokenizer_2
+# inpaintpipe.unet = pipe.unet
+# todo this should work
+# inpaintpipe = StableDiffusionXLInpaintPipeline( # construct an inpainter using the existing model
+#     vae=pipe.vae,
+#     text_encoder_2=pipe.text_encoder_2,
+#     text_encoder=pipe.text_encoder,
+#     unet=pipe.unet,
+#     scheduler=pipe.scheduler,
+#     tokenizer=pipe.tokenizer,
+#     tokenizer_2=pipe.tokenizer_2,
+#     requires_aesthetics_score=False,
+# )
+inpaintpipe.to("cuda")
+inpaintpipe.watermark = None
+# inpaintpipe.register_to_config(requires_aesthetics_score=False)
+
+inpaint_refiner = StableDiffusionXLInpaintPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-refiner-1.0",
+    text_encoder_2=inpaintpipe.text_encoder_2,
+    vae=inpaintpipe.vae,
+    torch_dtype=torch.bfloat16,
+    use_safetensors=True,
+    variant="fp16",
+
+    tokenizer_2=refiner.tokenizer_2,
+    tokenizer=refiner.tokenizer,
+    scheduler=refiner.scheduler,
+    text_encoder=refiner.text_encoder,
+    unet=refiner.unet,
+)
+# del inpaint_refiner.vae
+# del inpaint_refiner.text_encoder_2
+# del inpaint_refiner.text_encoder
+# del inpaint_refiner.scheduler
+# del inpaint_refiner.tokenizer
+# del inpaint_refiner.tokenizer_2
+# del inpaint_refiner.unet
+# inpaint_refiner.vae = inpaintpipe.vae
+# inpaint_refiner.text_encoder_2 = inpaintpipe.text_encoder_2
+#
+# inpaint_refiner.text_encoder = refiner.text_encoder
+# inpaint_refiner.scheduler = refiner.scheduler
+# inpaint_refiner.tokenizer = refiner.tokenizer
+# inpaint_refiner.tokenizer_2 = refiner.tokenizer_2
+# inpaint_refiner.unet = refiner.unet
+
+# inpaint_refiner = StableDiffusionXLInpaintPipeline(
+#     text_encoder_2=inpaintpipe.text_encoder_2,
+#     vae=inpaintpipe.vae,
+#     # the rest from the existing refiner
+#     tokenizer_2=refiner.tokenizer_2,
+#     tokenizer=refiner.tokenizer,
+#     scheduler=refiner.scheduler,
+#     text_encoder=refiner.text_encoder,
+#     unet=refiner.unet,
+#     requires_aesthetics_score=False,
+# )
+inpaint_refiner.to("cuda")
+inpaint_refiner.watermark = None
+# inpaint_refiner.register_to_config(requires_aesthetics_score=False)
+
+n_steps = 40
+high_noise_frac = 0.8
+
+# if using torch < 2.0
+# pipe.enable_xformers_memory_efficient_attention()
+
+
+# pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
+# this can cause errors on some inputs so consider disabling it
+pipe.unet = torch.compile(pipe.unet)
+refiner.unet = torch.compile(refiner.unet)#, mode="reduce-overhead", fullgraph=True)
+# compile the inpainters - todo reuse the other unets? swap out the models for others/del them so they share models and can be swapped efficiently
+inpaintpipe.unet = pipe.unet
+inpaint_refiner.unet = refiner.unet
+# inpaintpipe.unet = torch.compile(inpaintpipe.unet)
+# inpaint_refiner.unet = torch.compile(inpaint_refiner.unet)
+from pydantic import BaseModel
+
+app = FastAPI(
+    openapi_url="/static/openapi.json",
+    docs_url="/swagger-docs",
+    redoc_url="/redoc",
+    title="Generate Images Netwrck API",
+    description="Character Chat API",
+    # root_path="https://api.text-generator.io",
+    version="1",
+)
+app.add_middleware(GZipMiddleware, minimum_size=1000)
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+stopwords = nltk.corpus.stopwords.words("english")
+
+class Img(BaseModel):
+    system_prompt: str
+    ASSISTANT: str
+
+# img_url = "http://phlrr2019.guest.corp.microsoft.com:8000/img1_sdv2.1.png"
+img_url = "http://phlrr3058.guest.corp.microsoft.com:8000/"#/img1_sdv2.1.png"
+
+is_gpu_busy = False
+
+
+@app.post("/image_url")
+def image_url(img: Img):
+    system_prompt = img.system_prompt
+    prompt = img.ASSISTANT
+    # if Path(save_path).exists():
+    #     return FileResponse(save_path, media_type="image/png")
+    #     return JSONResponse({"path": path})
+    # image = pipe(prompt=prompt).images[0]
+    g = torch.Generator(device="cuda")
+    # image = pipe(prompt=prompt, width=1024, height=1024, generator=g).images[0]
+    image = pipe(prompt=prompt, width=1024, height=1024).images[0]
+
+    # if not save_path:
+    save_path = generate_save_path()
+    save_path = f"images/{save_path}.png"
+    image.save(save_path)
+    # save_path = '/'.join(path_components) + quote_plus(final_name)
+    path = f"{img_url}/{save_path}"
+    return JSONResponse({"path": path})
+
+
+@app.get("/make_image")
+# @app.post("/make_image")
+def make_image(prompt: str, save_path: str = ""):
+    if Path(save_path).exists():
+        return FileResponse(save_path, media_type="image/png")
+    image = pipe(prompt=prompt).images[0]
+    if not save_path:
+        save_path = f"images/{prompt}.png"
+    image.save(save_path)
+    return FileResponse(save_path, media_type="image/png")
+
+
+@app.get("/create_and_upload_image")
+def create_and_upload_image(prompt: str, width: int=1024, height:int=1024, save_path: str = ""):
+    path_components = save_path.split("/")[0:-1]
+    final_name = save_path.split("/")[-1]
+    if not path_components:
+        path_components = []
+    save_path = '/'.join(path_components) + quote_plus(final_name)
+    path = get_image_or_create_upload_to_cloud_storage(prompt, width, height, save_path)
+    return JSONResponse({"path": path})
+
+@app.get("/inpaint_and_upload_image")
+def inpaint_and_upload_image(prompt: str, image_url:str, mask_url:str, save_path: str = ""):
+    path_components = save_path.split("/")[0:-1]
+    final_name = save_path.split("/")[-1]
+    if not path_components:
+        path_components = []
+    save_path = '/'.join(path_components) + quote_plus(final_name)
+    path = get_image_or_inpaint_upload_to_cloud_storage(prompt, image_url, mask_url, save_path)
+    return JSONResponse({"path": path})
+
+
+def get_image_or_create_upload_to_cloud_storage(prompt:str,width:int, height:int, save_path:str):
+    prompt = shorten_too_long_text(prompt)
+    save_path = shorten_too_long_text(save_path)
+    # check exists - todo cache this
+    if check_if_blob_exists(save_path):
+        return f"https://{BUCKET_NAME}/{BUCKET_PATH}/{save_path}"
+    bio = create_image_from_prompt(prompt, width, height)
+    if bio is None:
+        return None # error thrown in pool
+    link = upload_to_bucket(save_path, bio, is_bytesio=True)
+    return link
+def get_image_or_inpaint_upload_to_cloud_storage(prompt:str, image_url:str, mask_url:str, save_path:str):
+    prompt = shorten_too_long_text(prompt)
+    save_path = shorten_too_long_text(save_path)
+    # check exists - todo cache this
+    if check_if_blob_exists(save_path):
+        return f"https://{BUCKET_NAME}/{BUCKET_PATH}/{save_path}"
+    bio = inpaint_image_from_prompt(prompt, image_url, mask_url)
+    if bio is None:
+        return None # error thrown in pool
+    link = upload_to_bucket(save_path, bio, is_bytesio=True)
+    return link
+
+# multiprocessing.set_start_method('spawn', True)
+# processes_pool = Pool(1) # cant do too much at once or OOM errors happen
+# def create_image_from_prompt_sync(prompt):
+#     """have to call this sync to avoid OOM errors"""
+#     return processes_pool.apply_async(create_image_from_prompt, args=(prompt,), ).wait()
+
+def create_image_from_prompt(prompt, width, height):
+    # round width and height down to multiple of 64
+    block_width = width - (width % 64)
+    block_height = height - (height % 64)
+    prompt = shorten_too_long_text(prompt)
+    # image = pipe(prompt=prompt).images[0]
+    try:
+        image = pipe(prompt=prompt,
+                     width=block_width,
+                     height=block_height,
+                     # denoising_end=high_noise_frac,
+                     # output_type='latent',
+                     # height=512,
+                     # width=512,
+                     num_inference_steps=50).images[0]  # normally uses 50 steps
+    except Exception as e:
+        # try rm stopwords + half the prompt
+        # todo try prompt permutations
+        logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+        prompt = ' '.join((word for word in prompt if word not in stopwords))
+        prompts = prompt.split()
+
+        prompt = ' '.join(prompts[:len(prompts) // 2])
+        logger.info(f"shortened prompt to: {len(prompt)}")
+        image = None
+        if prompt:
+            try:
+                image = pipe(prompt=prompt,
+                             width=block_width,
+                             height=block_height,
+                             # denoising_end=high_noise_frac,
+                             # output_type='latent',
+                             # height=512,
+                             # width=512,
+                             num_inference_steps=50).images[0]  # normally uses 50 steps
+            except Exception as e:
+                # logger.info("trying to permute prompt")
+                # # try two swaps of the prompt/permutations
+                # prompt = prompt.split()
+                # prompt = ' '.join(permutations(prompt, 2).__next__())
+                logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+                prompt = ' '.join((word for word in prompt if word not in stopwords))
+                prompts = prompt.split()
+
+                prompt = ' '.join(prompts[:len(prompts) // 2])
+                logger.info(f"shortened prompt to: {len(prompt)}")
+
+                try:
+                    image = pipe(prompt=prompt,
+                                 width=block_width,
+                                 height=block_height,
+                                 # denoising_end=high_noise_frac,
+                                 # output_type='latent', # dont need latent yet - we refine the image at full res
+                                 # height=512,
+                                 # width=512,
+                                 num_inference_steps=50).images[0]  # normally uses 50 steps
+                except Exception as e:
+                    # just error out
+                    traceback.print_exc()
+                    raise e
+                    # logger.info("restarting server to fix cuda issues (device side asserts)")
+                    # todo fix device side asserts instead of restart to fix
+                    # todo only restart the correct gunicorn
+                    # this could be really annoying if your running other gunicorns on your machine which also get restarted
+                    # os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+                    # os.system("kill -1 `pgrep gunicorn`")
+    # todo refine
+    # if image != None:
+    #     image = refiner(
+    #         prompt=prompt,
+    #         # width=block_width,
+    #         # height=block_height,
+    #         num_inference_steps=n_steps,
+    #         # denoising_start=high_noise_frac,
+    #         image=image,
+    #     ).images[0]
+    if width != block_width or height != block_height:
+        # resize to original size width/height
+        # find aspect ratio to scale up to that covers the original img input width/height
+        scale_up_ratio = max(width / block_width, height / block_height)
+        image = image.resize((math.ceil(block_width * scale_up_ratio), math.ceil(height * scale_up_ratio)))
+        # crop image to original size
+        image = image.crop((0, 0, width, height))
+    # try:
+    #     # gc.collect()
+    #     torch.cuda.empty_cache()
+    # except Exception as e:
+    #     traceback.print_exc()
+    #     logger.info("restarting server to fix cuda issues (device side asserts)")
+    #     # todo fix device side asserts instead of restart to fix
+    #     # todo only restart the correct gunicorn
+    #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+    #     os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+    #     os.system("kill -1 `pgrep gunicorn`")
+    # save as bytesio
+    bs = BytesIO()
+
+    bright_count = np.sum(np.array(image) > 0)
+    if bright_count == 0:
+        # we have a black image, this is an error likely we need a restart
+        logger.info("restarting server to fix cuda issues (device side asserts)")
+        #     # todo fix device side asserts instead of restart to fix
+        #     # todo only restart the correct gunicorn
+        #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+        os.system("kill -1 `pgrep gunicorn`")
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep uvicorn`")
+        os.system("kill -1 `pgrep uvicorn`")
+
+        return None
+    image.save(bs, quality=85, optimize=True, format="webp")
+    bio = bs.getvalue()
+    # touch progress.txt file - if we dont do this we get restarted by supervisor/other processes for reliability
+    with open("progress.txt", "w") as f:
+        current_time = datetime.now().strftime("%H:%M:%S")
+        f.write(f"{current_time}")
+    return bio
+
+def inpaint_image_from_prompt(prompt, image_url: str, mask_url: str):
+    prompt = shorten_too_long_text(prompt)
+    # image = pipe(prompt=prompt).images[0]
+
+    init_image = load_image(image_url).convert("RGB")
+    mask_image = load_image(mask_url).convert("RGB") # why rgb for a 1 channel mask?
+    num_inference_steps = 75
+    high_noise_frac = 0.7
+
+    try:
+        image = inpaintpipe(
+            prompt=prompt,
+            image=init_image,
+            mask_image=mask_image,
+            num_inference_steps=num_inference_steps,
+            denoising_start=high_noise_frac,
+            output_type="latent",
+        ).images[0]  # normally uses 50 steps
+    except Exception as e:
+        # try rm stopwords + half the prompt
+        # todo try prompt permutations
+        logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+        prompt = ' '.join((word for word in prompt if word not in stopwords))
+        prompts = prompt.split()
+
+        prompt = ' '.join(prompts[:len(prompts) // 2])
+        logger.info(f"shortened prompt to: {len(prompt)}")
+        image = None
+        if prompt:
+            try:
+                image = pipe(
+                    prompt=prompt,
+                    image=init_image,
+                    mask_image=mask_image,
+                    num_inference_steps=num_inference_steps,
+                    denoising_start=high_noise_frac,
+                    output_type="latent",
+                ).images[0]  # normally uses 50 steps
+            except Exception as e:
+                # logger.info("trying to permute prompt")
+                # # try two swaps of the prompt/permutations
+                # prompt = prompt.split()
+                # prompt = ' '.join(permutations(prompt, 2).__next__())
+                logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+                prompt = ' '.join((word for word in prompt if word not in stopwords))
+                prompts = prompt.split()
+
+                prompt = ' '.join(prompts[:len(prompts) // 2])
+                logger.info(f"shortened prompt to: {len(prompt)}")
+
+                try:
+                    image = inpaintpipe(
+                        prompt=prompt,
+                        image=init_image,
+                        mask_image=mask_image,
+                        num_inference_steps=num_inference_steps,
+                        denoising_start=high_noise_frac,
+                        output_type="latent",
+                    ).images[0]  # normally uses 50 steps
+                except Exception as e:
+                    # just error out
+                    traceback.print_exc()
+                    raise e
+                    # logger.info("restarting server to fix cuda issues (device side asserts)")
+                    # todo fix device side asserts instead of restart to fix
+                    # todo only restart the correct gunicorn
+                    # this could be really annoying if your running other gunicorns on your machine which also get restarted
+                    # os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+                    # os.system("kill -1 `pgrep gunicorn`")
+    if image != None:
+        image = inpaint_refiner(
+            prompt=prompt,
+            image=image,
+            mask_image=mask_image,
+            num_inference_steps=num_inference_steps,
+            denoising_start=high_noise_frac,
+
+        ).images[0]
+    # try:
+    #     # gc.collect()
+    #     torch.cuda.empty_cache()
+    # except Exception as e:
+    #     traceback.print_exc()
+    #     logger.info("restarting server to fix cuda issues (device side asserts)")
+    #     # todo fix device side asserts instead of restart to fix
+    #     # todo only restart the correct gunicorn
+    #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+    #     os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+    #     os.system("kill -1 `pgrep gunicorn`")
+    # save as bytesio
+    bs = BytesIO()
+
+    bright_count = np.sum(np.array(image) > 0)
+    if bright_count == 0:
+        # we have a black image, this is an error likely we need a restart
+        logger.info("restarting server to fix cuda issues (device side asserts)")
+        #     # todo fix device side asserts instead of restart to fix
+        #     # todo only restart the correct gunicorn
+        #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+        os.system("kill -1 `pgrep gunicorn`")
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep uvicorn`")
+        os.system("kill -1 `pgrep uvicorn`")
+
+        return None
+    image.save(bs, quality=85, optimize=True, format="webp")
+    bio = bs.getvalue()
+    # touch progress.txt file - if we dont do this we get restarted by supervisor/other processes for reliability
+    with open("progress.txt", "w") as f:
+        current_time = datetime.now().strftime("%H:%M:%S")
+        f.write(f"{current_time}")
+    return bio
+
+
+
+def shorten_too_long_text(prompt):
+    if len(prompt) > 200:
+        # remove stopwords
+        prompt = prompt.split() # todo also split hyphens
+        prompt = ' '.join((word for word in prompt if word not in stopwords))
+        if len(prompt) > 200:
+            prompt = prompt[:200]
+    return prompt
+
+# image = pipe(prompt=prompt).images[0]
+#
+# image.save("test.png")
+# # save all images
+# for i, image in enumerate(images):
+#     image.save(f"{i}.png")
+

img/stable-diffusion-server/main_v2.py

@@ -0,0 +1,548 @@
+import gc
+import math
+import multiprocessing
+import os
+import traceback
+from datetime import datetime
+from io import BytesIO
+from itertools import permutations
+from multiprocessing.pool import Pool
+from pathlib import Path
+from urllib.parse import quote_plus
+
+import numpy as np
+import nltk
+import torch
+
+from PIL.Image import Image
+from diffusers import DiffusionPipeline, StableDiffusionXLInpaintPipeline
+from diffusers.utils import load_image
+from fastapi import FastAPI
+from fastapi.middleware.gzip import GZipMiddleware
+from loguru import logger
+from starlette.middleware.cors import CORSMiddleware
+from starlette.responses import FileResponse
+from starlette.responses import JSONResponse
+
+from env import BUCKET_PATH, BUCKET_NAME
+# from stable_diffusion_server.bucket_api import check_if_blob_exists, upload_to_bucket
+torch._dynamo.config.suppress_errors = True
+
+import string
+import random
+
+def generate_save_path():
+    # initializing size of string
+    N = 7
+
+    # using random.choices()
+    # generating random strings
+    res = ''.join(random.choices(string.ascii_uppercase +
+                                 string.digits, k=N))
+    return res
+
+# pipe = DiffusionPipeline.from_pretrained(
+#     "models/stable-diffusion-xl-base-1.0",
+#     torch_dtype=torch.bfloat16,
+#     use_safetensors=True,
+#     variant="fp16",
+#     # safety_checker=None,
+# )  # todo try torch_dtype=bfloat16
+
+model_dir = os.getenv("SDXL_MODEL_DIR")
+
+if model_dir:
+    # Use local model
+    model_key_base = os.path.join(model_dir, "stable-diffusion-xl-base-1.0")
+    model_key_refiner = os.path.join(model_dir, "stable-diffusion-xl-refiner-1.0")
+else:
+    model_key_base = "stabilityai/stable-diffusion-xl-base-1.0"
+    model_key_refiner = "stabilityai/stable-diffusion-xl-refiner-1.0"
+
+pipe = DiffusionPipeline.from_pretrained(model_key_base, torch_dtype=torch.float16, use_safetensors=True, variant="fp16")
+
+pipe.watermark = None
+
+pipe.to("cuda")
+
+refiner = DiffusionPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-refiner-1.0",
+    text_encoder_2=pipe.text_encoder_2,
+    vae=pipe.vae,
+    torch_dtype=torch.bfloat16, # safer to use bfloat?
+    use_safetensors=True,
+    variant="fp16", #remember not to download the big model
+)
+refiner.watermark = None
+refiner.to("cuda")
+
+# {'scheduler', 'text_encoder', 'text_encoder_2', 'tokenizer', 'tokenizer_2', 'unet', 'vae'} can be passed in from existing model
+inpaintpipe = StableDiffusionXLInpaintPipeline.from_pretrained(
+    "models/stable-diffusion-xl-base-1.0", torch_dtype=torch.bfloat16, variant="fp16", use_safetensors=True,
+    scheduler=pipe.scheduler,
+    text_encoder=pipe.text_encoder,
+    text_encoder_2=pipe.text_encoder_2,
+    tokenizer=pipe.tokenizer,
+    tokenizer_2=pipe.tokenizer_2,
+    unet=pipe.unet,
+    vae=pipe.vae,
+    # load_connected_pipeline=
+)
+# # switch out to save gpu mem
+# del inpaintpipe.vae
+# del inpaintpipe.text_encoder_2
+# del inpaintpipe.text_encoder
+# del inpaintpipe.scheduler
+# del inpaintpipe.tokenizer
+# del inpaintpipe.tokenizer_2
+# del inpaintpipe.unet
+# inpaintpipe.vae = pipe.vae
+# inpaintpipe.text_encoder_2 = pipe.text_encoder_2
+# inpaintpipe.text_encoder = pipe.text_encoder
+# inpaintpipe.scheduler = pipe.scheduler
+# inpaintpipe.tokenizer = pipe.tokenizer
+# inpaintpipe.tokenizer_2 = pipe.tokenizer_2
+# inpaintpipe.unet = pipe.unet
+# todo this should work
+# inpaintpipe = StableDiffusionXLInpaintPipeline( # construct an inpainter using the existing model
+#     vae=pipe.vae,
+#     text_encoder_2=pipe.text_encoder_2,
+#     text_encoder=pipe.text_encoder,
+#     unet=pipe.unet,
+#     scheduler=pipe.scheduler,
+#     tokenizer=pipe.tokenizer,
+#     tokenizer_2=pipe.tokenizer_2,
+#     requires_aesthetics_score=False,
+# )
+inpaintpipe.to("cuda")
+inpaintpipe.watermark = None
+# inpaintpipe.register_to_config(requires_aesthetics_score=False)
+
+inpaint_refiner = StableDiffusionXLInpaintPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-refiner-1.0",
+    text_encoder_2=inpaintpipe.text_encoder_2,
+    vae=inpaintpipe.vae,
+    torch_dtype=torch.bfloat16,
+    use_safetensors=True,
+    variant="fp16",
+
+    tokenizer_2=refiner.tokenizer_2,
+    tokenizer=refiner.tokenizer,
+    scheduler=refiner.scheduler,
+    text_encoder=refiner.text_encoder,
+    unet=refiner.unet,
+)
+# del inpaint_refiner.vae
+# del inpaint_refiner.text_encoder_2
+# del inpaint_refiner.text_encoder
+# del inpaint_refiner.scheduler
+# del inpaint_refiner.tokenizer
+# del inpaint_refiner.tokenizer_2
+# del inpaint_refiner.unet
+# inpaint_refiner.vae = inpaintpipe.vae
+# inpaint_refiner.text_encoder_2 = inpaintpipe.text_encoder_2
+#
+# inpaint_refiner.text_encoder = refiner.text_encoder
+# inpaint_refiner.scheduler = refiner.scheduler
+# inpaint_refiner.tokenizer = refiner.tokenizer
+# inpaint_refiner.tokenizer_2 = refiner.tokenizer_2
+# inpaint_refiner.unet = refiner.unet
+
+# inpaint_refiner = StableDiffusionXLInpaintPipeline(
+#     text_encoder_2=inpaintpipe.text_encoder_2,
+#     vae=inpaintpipe.vae,
+#     # the rest from the existing refiner
+#     tokenizer_2=refiner.tokenizer_2,
+#     tokenizer=refiner.tokenizer,
+#     scheduler=refiner.scheduler,
+#     text_encoder=refiner.text_encoder,
+#     unet=refiner.unet,
+#     requires_aesthetics_score=False,
+# )
+inpaint_refiner.to("cuda")
+inpaint_refiner.watermark = None
+# inpaint_refiner.register_to_config(requires_aesthetics_score=False)
+
+n_steps = 40
+high_noise_frac = 0.8
+
+# if using torch < 2.0
+# pipe.enable_xformers_memory_efficient_attention()
+
+
+# pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
+# this can cause errors on some inputs so consider disabling it
+pipe.unet = torch.compile(pipe.unet)
+refiner.unet = torch.compile(refiner.unet)#, mode="reduce-overhead", fullgraph=True)
+# compile the inpainters - todo reuse the other unets? swap out the models for others/del them so they share models and can be swapped efficiently
+inpaintpipe.unet = pipe.unet
+inpaint_refiner.unet = refiner.unet
+# inpaintpipe.unet = torch.compile(inpaintpipe.unet)
+# inpaint_refiner.unet = torch.compile(inpaint_refiner.unet)
+from pydantic import BaseModel
+
+app = FastAPI(
+    openapi_url="/static/openapi.json",
+    docs_url="/swagger-docs",
+    redoc_url="/redoc",
+    title="Generate Images Netwrck API",
+    description="Character Chat API",
+    # root_path="https://api.text-generator.io",
+    version="1",
+)
+app.add_middleware(GZipMiddleware, minimum_size=1000)
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+stopwords = nltk.corpus.stopwords.words("english")
+
+class Img(BaseModel):
+    system_prompt: str
+    ASSISTANT: str
+
+# img_url = "http://phlrr2019.guest.corp.microsoft.com:8000/img1_sdv2.1.png"
+img_url = "http://phlrr3105.guest.corp.microsoft.com:8000/"#/img1_sdv2.1.png"
+
+is_gpu_busy = False
+
+
+@app.post("/image_url")
+def image_url(img: Img):
+    system_prompt = img.system_prompt
+    prompt = img.ASSISTANT
+    # if Path(save_path).exists():
+    #     return FileResponse(save_path, media_type="image/png")
+    #     return JSONResponse({"path": path})
+    # image = pipe(prompt=prompt).images[0]
+    g = torch.Generator(device="cuda")
+    image = pipe(prompt=prompt, width=1024, height=1024, generator=g).images[0]
+
+    # if not save_path:
+    save_path = generate_save_path()
+    save_path = f"images/{save_path}.png"
+    image.save(save_path)
+    # save_path = '/'.join(path_components) + quote_plus(final_name)
+    path = f"{img_url}/{save_path}"
+    return JSONResponse({"path": path})
+
+
+@app.get("/make_image")
+# @app.post("/make_image")
+def make_image(prompt: str, save_path: str = ""):
+    if Path(save_path).exists():
+        return FileResponse(save_path, media_type="image/png")
+    image = pipe(prompt=prompt).images[0]
+    if not save_path:
+        save_path = f"images/{prompt}.png"
+    image.save(save_path)
+    return FileResponse(save_path, media_type="image/png")
+
+
+@app.get("/create_and_upload_image")
+def create_and_upload_image(prompt: str, width: int=1024, height:int=1024, save_path: str = ""):
+    path_components = save_path.split("/")[0:-1]
+    final_name = save_path.split("/")[-1]
+    if not path_components:
+        path_components = []
+    save_path = '/'.join(path_components) + quote_plus(final_name)
+    path = get_image_or_create_upload_to_cloud_storage(prompt, width, height, save_path)
+    return JSONResponse({"path": path})
+
+@app.get("/inpaint_and_upload_image")
+def inpaint_and_upload_image(prompt: str, image_url:str, mask_url:str, save_path: str = ""):
+    path_components = save_path.split("/")[0:-1]
+    final_name = save_path.split("/")[-1]
+    if not path_components:
+        path_components = []
+    save_path = '/'.join(path_components) + quote_plus(final_name)
+    path = get_image_or_inpaint_upload_to_cloud_storage(prompt, image_url, mask_url, save_path)
+    return JSONResponse({"path": path})
+
+
+def get_image_or_create_upload_to_cloud_storage(prompt:str,width:int, height:int, save_path:str):
+    prompt = shorten_too_long_text(prompt)
+    save_path = shorten_too_long_text(save_path)
+    # check exists - todo cache this
+    if check_if_blob_exists(save_path):
+        return f"https://{BUCKET_NAME}/{BUCKET_PATH}/{save_path}"
+    bio = create_image_from_prompt(prompt, width, height)
+    if bio is None:
+        return None # error thrown in pool
+    link = upload_to_bucket(save_path, bio, is_bytesio=True)
+    return link
+def get_image_or_inpaint_upload_to_cloud_storage(prompt:str, image_url:str, mask_url:str, save_path:str):
+    prompt = shorten_too_long_text(prompt)
+    save_path = shorten_too_long_text(save_path)
+    # check exists - todo cache this
+    if check_if_blob_exists(save_path):
+        return f"https://{BUCKET_NAME}/{BUCKET_PATH}/{save_path}"
+    bio = inpaint_image_from_prompt(prompt, image_url, mask_url)
+    if bio is None:
+        return None # error thrown in pool
+    link = upload_to_bucket(save_path, bio, is_bytesio=True)
+    return link
+
+# multiprocessing.set_start_method('spawn', True)
+# processes_pool = Pool(1) # cant do too much at once or OOM errors happen
+# def create_image_from_prompt_sync(prompt):
+#     """have to call this sync to avoid OOM errors"""
+#     return processes_pool.apply_async(create_image_from_prompt, args=(prompt,), ).wait()
+
+def create_image_from_prompt(prompt, width, height):
+    # round width and height down to multiple of 64
+    block_width = width - (width % 64)
+    block_height = height - (height % 64)
+    prompt = shorten_too_long_text(prompt)
+    # image = pipe(prompt=prompt).images[0]
+    try:
+        image = pipe(prompt=prompt,
+                     width=block_width,
+                     height=block_height,
+                     # denoising_end=high_noise_frac,
+                     # output_type='latent',
+                     # height=512,
+                     # width=512,
+                     num_inference_steps=50).images[0]  # normally uses 50 steps
+    except Exception as e:
+        # try rm stopwords + half the prompt
+        # todo try prompt permutations
+        logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+        prompt = ' '.join((word for word in prompt if word not in stopwords))
+        prompts = prompt.split()
+
+        prompt = ' '.join(prompts[:len(prompts) // 2])
+        logger.info(f"shortened prompt to: {len(prompt)}")
+        image = None
+        if prompt:
+            try:
+                image = pipe(prompt=prompt,
+                             width=block_width,
+                             height=block_height,
+                             # denoising_end=high_noise_frac,
+                             # output_type='latent',
+                             # height=512,
+                             # width=512,
+                             num_inference_steps=50).images[0]  # normally uses 50 steps
+            except Exception as e:
+                # logger.info("trying to permute prompt")
+                # # try two swaps of the prompt/permutations
+                # prompt = prompt.split()
+                # prompt = ' '.join(permutations(prompt, 2).__next__())
+                logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+                prompt = ' '.join((word for word in prompt if word not in stopwords))
+                prompts = prompt.split()
+
+                prompt = ' '.join(prompts[:len(prompts) // 2])
+                logger.info(f"shortened prompt to: {len(prompt)}")
+
+                try:
+                    image = pipe(prompt=prompt,
+                                 width=block_width,
+                                 height=block_height,
+                                 # denoising_end=high_noise_frac,
+                                 # output_type='latent', # dont need latent yet - we refine the image at full res
+                                 # height=512,
+                                 # width=512,
+                                 num_inference_steps=50).images[0]  # normally uses 50 steps
+                except Exception as e:
+                    # just error out
+                    traceback.print_exc()
+                    raise e
+                    # logger.info("restarting server to fix cuda issues (device side asserts)")
+                    # todo fix device side asserts instead of restart to fix
+                    # todo only restart the correct gunicorn
+                    # this could be really annoying if your running other gunicorns on your machine which also get restarted
+                    # os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+                    # os.system("kill -1 `pgrep gunicorn`")
+    # todo refine
+    # if image != None:
+    #     image = refiner(
+    #         prompt=prompt,
+    #         # width=block_width,
+    #         # height=block_height,
+    #         num_inference_steps=n_steps,
+    #         # denoising_start=high_noise_frac,
+    #         image=image,
+    #     ).images[0]
+    if width != block_width or height != block_height:
+        # resize to original size width/height
+        # find aspect ratio to scale up to that covers the original img input width/height
+        scale_up_ratio = max(width / block_width, height / block_height)
+        image = image.resize((math.ceil(block_width * scale_up_ratio), math.ceil(height * scale_up_ratio)))
+        # crop image to original size
+        image = image.crop((0, 0, width, height))
+    # try:
+    #     # gc.collect()
+    #     torch.cuda.empty_cache()
+    # except Exception as e:
+    #     traceback.print_exc()
+    #     logger.info("restarting server to fix cuda issues (device side asserts)")
+    #     # todo fix device side asserts instead of restart to fix
+    #     # todo only restart the correct gunicorn
+    #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+    #     os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+    #     os.system("kill -1 `pgrep gunicorn`")
+    # save as bytesio
+    bs = BytesIO()
+
+    bright_count = np.sum(np.array(image) > 0)
+    if bright_count == 0:
+        # we have a black image, this is an error likely we need a restart
+        logger.info("restarting server to fix cuda issues (device side asserts)")
+        #     # todo fix device side asserts instead of restart to fix
+        #     # todo only restart the correct gunicorn
+        #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+        os.system("kill -1 `pgrep gunicorn`")
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep uvicorn`")
+        os.system("kill -1 `pgrep uvicorn`")
+
+        return None
+    image.save(bs, quality=85, optimize=True, format="webp")
+    bio = bs.getvalue()
+    # touch progress.txt file - if we dont do this we get restarted by supervisor/other processes for reliability
+    with open("progress.txt", "w") as f:
+        current_time = datetime.now().strftime("%H:%M:%S")
+        f.write(f"{current_time}")
+    return bio
+
+def inpaint_image_from_prompt(prompt, image_url: str, mask_url: str):
+    prompt = shorten_too_long_text(prompt)
+    # image = pipe(prompt=prompt).images[0]
+
+    init_image = load_image(image_url).convert("RGB")
+    mask_image = load_image(mask_url).convert("RGB") # why rgb for a 1 channel mask?
+    num_inference_steps = 75
+    high_noise_frac = 0.7
+
+    try:
+        image = inpaintpipe(
+            prompt=prompt,
+            image=init_image,
+            mask_image=mask_image,
+            num_inference_steps=num_inference_steps,
+            denoising_start=high_noise_frac,
+            output_type="latent",
+        ).images[0]  # normally uses 50 steps
+    except Exception as e:
+        # try rm stopwords + half the prompt
+        # todo try prompt permutations
+        logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+        prompt = ' '.join((word for word in prompt if word not in stopwords))
+        prompts = prompt.split()
+
+        prompt = ' '.join(prompts[:len(prompts) // 2])
+        logger.info(f"shortened prompt to: {len(prompt)}")
+        image = None
+        if prompt:
+            try:
+                image = pipe(
+                    prompt=prompt,
+                    image=init_image,
+                    mask_image=mask_image,
+                    num_inference_steps=num_inference_steps,
+                    denoising_start=high_noise_frac,
+                    output_type="latent",
+                ).images[0]  # normally uses 50 steps
+            except Exception as e:
+                # logger.info("trying to permute prompt")
+                # # try two swaps of the prompt/permutations
+                # prompt = prompt.split()
+                # prompt = ' '.join(permutations(prompt, 2).__next__())
+                logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+                prompt = ' '.join((word for word in prompt if word not in stopwords))
+                prompts = prompt.split()
+
+                prompt = ' '.join(prompts[:len(prompts) // 2])
+                logger.info(f"shortened prompt to: {len(prompt)}")
+
+                try:
+                    image = inpaintpipe(
+                        prompt=prompt,
+                        image=init_image,
+                        mask_image=mask_image,
+                        num_inference_steps=num_inference_steps,
+                        denoising_start=high_noise_frac,
+                        output_type="latent",
+                    ).images[0]  # normally uses 50 steps
+                except Exception as e:
+                    # just error out
+                    traceback.print_exc()
+                    raise e
+                    # logger.info("restarting server to fix cuda issues (device side asserts)")
+                    # todo fix device side asserts instead of restart to fix
+                    # todo only restart the correct gunicorn
+                    # this could be really annoying if your running other gunicorns on your machine which also get restarted
+                    # os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+                    # os.system("kill -1 `pgrep gunicorn`")
+    if image != None:
+        image = inpaint_refiner(
+            prompt=prompt,
+            image=image,
+            mask_image=mask_image,
+            num_inference_steps=num_inference_steps,
+            denoising_start=high_noise_frac,
+
+        ).images[0]
+    # try:
+    #     # gc.collect()
+    #     torch.cuda.empty_cache()
+    # except Exception as e:
+    #     traceback.print_exc()
+    #     logger.info("restarting server to fix cuda issues (device side asserts)")
+    #     # todo fix device side asserts instead of restart to fix
+    #     # todo only restart the correct gunicorn
+    #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+    #     os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+    #     os.system("kill -1 `pgrep gunicorn`")
+    # save as bytesio
+    bs = BytesIO()
+
+    bright_count = np.sum(np.array(image) > 0)
+    if bright_count == 0:
+        # we have a black image, this is an error likely we need a restart
+        logger.info("restarting server to fix cuda issues (device side asserts)")
+        #     # todo fix device side asserts instead of restart to fix
+        #     # todo only restart the correct gunicorn
+        #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+        os.system("kill -1 `pgrep gunicorn`")
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep uvicorn`")
+        os.system("kill -1 `pgrep uvicorn`")
+
+        return None
+    image.save(bs, quality=85, optimize=True, format="webp")
+    bio = bs.getvalue()
+    # touch progress.txt file - if we dont do this we get restarted by supervisor/other processes for reliability
+    with open("progress.txt", "w") as f:
+        current_time = datetime.now().strftime("%H:%M:%S")
+        f.write(f"{current_time}")
+    return bio
+
+
+
+def shorten_too_long_text(prompt):
+    if len(prompt) > 200:
+        # remove stopwords
+        prompt = prompt.split() # todo also split hyphens
+        prompt = ' '.join((word for word in prompt if word not in stopwords))
+        if len(prompt) > 200:
+            prompt = prompt[:200]
+    return prompt
+
+# image = pipe(prompt=prompt).images[0]
+#
+# image.save("test.png")
+# # save all images
+# for i, image in enumerate(images):
+#     image.save(f"{i}.png")
+

img/stable-diffusion-server/main_v3.py

@@ -0,0 +1,578 @@
+import gc
+import math
+import multiprocessing
+import os
+import traceback
+from datetime import datetime
+from io import BytesIO
+from itertools import permutations
+from multiprocessing.pool import Pool
+from pathlib import Path
+from urllib.parse import quote_plus
+
+import numpy as np
+import nltk
+import torch
+
+from PIL.Image import Image
+from diffusers import DiffusionPipeline, StableDiffusionXLInpaintPipeline
+from diffusers.utils import load_image
+from fastapi import FastAPI
+from fastapi.middleware.gzip import GZipMiddleware
+from loguru import logger
+from starlette.middleware.cors import CORSMiddleware
+from starlette.responses import FileResponse
+from starlette.responses import JSONResponse
+
+from env import BUCKET_PATH, BUCKET_NAME
+# from stable_diffusion_server.bucket_api import check_if_blob_exists, upload_to_bucket
+torch._dynamo.config.suppress_errors = True
+
+import string
+import random
+
+def generate_save_path():
+    # initializing size of string
+    N = 7
+
+    # using random.choices()
+    # generating random strings
+    res = ''.join(random.choices(string.ascii_uppercase +
+                                 string.digits, k=N))
+    return res
+
+# pipe = DiffusionPipeline.from_pretrained(
+#     "models/stable-diffusion-xl-base-1.0",
+#     torch_dtype=torch.bfloat16,
+#     use_safetensors=True,
+#     variant="fp16",
+#     # safety_checker=None,
+# )  # todo try torch_dtype=bfloat16
+
+model_dir = os.getenv("SDXL_MODEL_DIR")
+
+if model_dir:
+    # Use local model
+    model_key_base = os.path.join(model_dir, "stable-diffusion-xl-base-1.0")
+    model_key_refiner = os.path.join(model_dir, "stable-diffusion-xl-refiner-1.0")
+else:
+    model_key_base = "stabilityai/stable-diffusion-xl-base-1.0"
+    model_key_refiner = "stabilityai/stable-diffusion-xl-refiner-1.0"
+
+pipe = DiffusionPipeline.from_pretrained(model_key_base, torch_dtype=torch.float16, use_safetensors=True, variant="fp16")
+
+pipe.watermark = None
+
+pipe.to("cuda")
+
+refiner = DiffusionPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-refiner-1.0",
+    text_encoder_2=pipe.text_encoder_2,
+    vae=pipe.vae,
+    torch_dtype=torch.bfloat16, # safer to use bfloat?
+    use_safetensors=True,
+    variant="fp16", #remember not to download the big model
+)
+refiner.watermark = None
+refiner.to("cuda")
+
+# {'scheduler', 'text_encoder', 'text_encoder_2', 'tokenizer', 'tokenizer_2', 'unet', 'vae'} can be passed in from existing model
+inpaintpipe = StableDiffusionXLInpaintPipeline.from_pretrained(
+    "models/stable-diffusion-xl-base-1.0", torch_dtype=torch.bfloat16, variant="fp16", use_safetensors=True,
+    scheduler=pipe.scheduler,
+    text_encoder=pipe.text_encoder,
+    text_encoder_2=pipe.text_encoder_2,
+    tokenizer=pipe.tokenizer,
+    tokenizer_2=pipe.tokenizer_2,
+    unet=pipe.unet,
+    vae=pipe.vae,
+    # load_connected_pipeline=
+)
+# # switch out to save gpu mem
+# del inpaintpipe.vae
+# del inpaintpipe.text_encoder_2
+# del inpaintpipe.text_encoder
+# del inpaintpipe.scheduler
+# del inpaintpipe.tokenizer
+# del inpaintpipe.tokenizer_2
+# del inpaintpipe.unet
+# inpaintpipe.vae = pipe.vae
+# inpaintpipe.text_encoder_2 = pipe.text_encoder_2
+# inpaintpipe.text_encoder = pipe.text_encoder
+# inpaintpipe.scheduler = pipe.scheduler
+# inpaintpipe.tokenizer = pipe.tokenizer
+# inpaintpipe.tokenizer_2 = pipe.tokenizer_2
+# inpaintpipe.unet = pipe.unet
+# todo this should work
+# inpaintpipe = StableDiffusionXLInpaintPipeline( # construct an inpainter using the existing model
+#     vae=pipe.vae,
+#     text_encoder_2=pipe.text_encoder_2,
+#     text_encoder=pipe.text_encoder,
+#     unet=pipe.unet,
+#     scheduler=pipe.scheduler,
+#     tokenizer=pipe.tokenizer,
+#     tokenizer_2=pipe.tokenizer_2,
+#     requires_aesthetics_score=False,
+# )
+inpaintpipe.to("cuda")
+inpaintpipe.watermark = None
+# inpaintpipe.register_to_config(requires_aesthetics_score=False)
+
+inpaint_refiner = StableDiffusionXLInpaintPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-refiner-1.0",
+    text_encoder_2=inpaintpipe.text_encoder_2,
+    vae=inpaintpipe.vae,
+    torch_dtype=torch.bfloat16,
+    use_safetensors=True,
+    variant="fp16",
+
+    tokenizer_2=refiner.tokenizer_2,
+    tokenizer=refiner.tokenizer,
+    scheduler=refiner.scheduler,
+    text_encoder=refiner.text_encoder,
+    unet=refiner.unet,
+)
+# del inpaint_refiner.vae
+# del inpaint_refiner.text_encoder_2
+# del inpaint_refiner.text_encoder
+# del inpaint_refiner.scheduler
+# del inpaint_refiner.tokenizer
+# del inpaint_refiner.tokenizer_2
+# del inpaint_refiner.unet
+# inpaint_refiner.vae = inpaintpipe.vae
+# inpaint_refiner.text_encoder_2 = inpaintpipe.text_encoder_2
+#
+# inpaint_refiner.text_encoder = refiner.text_encoder
+# inpaint_refiner.scheduler = refiner.scheduler
+# inpaint_refiner.tokenizer = refiner.tokenizer
+# inpaint_refiner.tokenizer_2 = refiner.tokenizer_2
+# inpaint_refiner.unet = refiner.unet
+
+# inpaint_refiner = StableDiffusionXLInpaintPipeline(
+#     text_encoder_2=inpaintpipe.text_encoder_2,
+#     vae=inpaintpipe.vae,
+#     # the rest from the existing refiner
+#     tokenizer_2=refiner.tokenizer_2,
+#     tokenizer=refiner.tokenizer,
+#     scheduler=refiner.scheduler,
+#     text_encoder=refiner.text_encoder,
+#     unet=refiner.unet,
+#     requires_aesthetics_score=False,
+# )
+inpaint_refiner.to("cuda")
+inpaint_refiner.watermark = None
+# inpaint_refiner.register_to_config(requires_aesthetics_score=False)
+
+n_steps = 40
+high_noise_frac = 0.8
+
+# if using torch < 2.0
+# pipe.enable_xformers_memory_efficient_attention()
+
+
+# pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
+# this can cause errors on some inputs so consider disabling it
+pipe.unet = torch.compile(pipe.unet)
+refiner.unet = torch.compile(refiner.unet)#, mode="reduce-overhead", fullgraph=True)
+# compile the inpainters - todo reuse the other unets? swap out the models for others/del them so they share models and can be swapped efficiently
+inpaintpipe.unet = pipe.unet
+inpaint_refiner.unet = refiner.unet
+# inpaintpipe.unet = torch.compile(inpaintpipe.unet)
+# inpaint_refiner.unet = torch.compile(inpaint_refiner.unet)
+from pydantic import BaseModel
+
+app = FastAPI(
+    openapi_url="/static/openapi.json",
+    docs_url="/swagger-docs",
+    redoc_url="/redoc",
+    title="Generate Images Netwrck API",
+    description="Character Chat API",
+    # root_path="https://api.text-generator.io",
+    version="1",
+)
+app.add_middleware(GZipMiddleware, minimum_size=1000)
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+stopwords = nltk.corpus.stopwords.words("english")
+
+class Img(BaseModel):
+    system_prompt: str
+    ASSISTANT: str
+
+# img_url = "http://phlrr2019.guest.corp.microsoft.com:8000/img1_sdv2.1.png"
+img_url = "http://phlrr3105.guest.corp.microsoft.com:8000/"#/img1_sdv2.1.png"
+
+is_gpu_busy = False
+
+def get_summary(system_prompt, prompt):
+    import requests
+    import time
+    from io import BytesIO
+    import json
+    summary_sys = """I want you to act as a text summarizer to help me create a concise summary of the text I provide. The summary can be up to 60.0 words in length, expressing the key points, key scenarios, main character and concepts written in the original text without adding your interpretations."""
+    instruction = summary_sys
+    # for human, assistant in history:
+    #     instruction += 'USER: ' + human + ' ASSISTANT: ' + assistant + '</s>'
+    # prompt = system_prompt + prompt
+    message = f"""My first request is to summarize this text – [{prompt}]"""
+    instruction += ' USER: ' + message + ' ASSISTANT:'
+
+    print("Ins: ", instruction)
+    # generate_response = requests.post("http://10.185.12.207:4455/stable_diffusion", json={"prompt": prompt})
+    # prompt = f""" My first request is to summarize this text – [{prompt}]"""
+    json_object = {"prompt": instruction,
+                   # "max_tokens": 2048000,
+                   "max_tokens": 90,
+                   "n": 1
+                   }
+    generate_response = requests.post("http://phlrr3105.guest.corp.microsoft.com:7991/generate", json=json_object)
+    # print(generate_response.content)
+    res_json = json.loads(generate_response.content)
+    ASSISTANT = res_json['text'][-1].split("ASSISTANT:")[-1].strip()
+    print(ASSISTANT)
+    return ASSISTANT
+
+@app.post("/image_url")
+def image_url(img: Img):
+    system_prompt = img.system_prompt
+    prompt = img.ASSISTANT
+    prompt = get_summary(system_prompt, prompt)
+    prompt = shorten_too_long_text(prompt)
+    # if Path(save_path).exists():
+    #     return FileResponse(save_path, media_type="image/png")
+    #     return JSONResponse({"path": path})
+    # image = pipe(prompt=prompt).images[0]
+    g = torch.Generator(device="cuda")
+    image = pipe(prompt=prompt, width=1024, height=1024, generator=g).images[0]
+
+    # if not save_path:
+    save_path = generate_save_path()
+    save_path = f"images/{save_path}.png"
+    image.save(save_path)
+    # save_path = '/'.join(path_components) + quote_plus(final_name)
+    path = f"{img_url}/{save_path}"
+    return JSONResponse({"path": path})
+
+
+@app.get("/make_image")
+# @app.post("/make_image")
+def make_image(prompt: str, save_path: str = ""):
+    if Path(save_path).exists():
+        return FileResponse(save_path, media_type="image/png")
+    image = pipe(prompt=prompt).images[0]
+    if not save_path:
+        save_path = f"images/{prompt}.png"
+    image.save(save_path)
+    return FileResponse(save_path, media_type="image/png")
+
+
+@app.get("/create_and_upload_image")
+def create_and_upload_image(prompt: str, width: int=1024, height:int=1024, save_path: str = ""):
+    path_components = save_path.split("/")[0:-1]
+    final_name = save_path.split("/")[-1]
+    if not path_components:
+        path_components = []
+    save_path = '/'.join(path_components) + quote_plus(final_name)
+    path = get_image_or_create_upload_to_cloud_storage(prompt, width, height, save_path)
+    return JSONResponse({"path": path})
+
+@app.get("/inpaint_and_upload_image")
+def inpaint_and_upload_image(prompt: str, image_url:str, mask_url:str, save_path: str = ""):
+    path_components = save_path.split("/")[0:-1]
+    final_name = save_path.split("/")[-1]
+    if not path_components:
+        path_components = []
+    save_path = '/'.join(path_components) + quote_plus(final_name)
+    path = get_image_or_inpaint_upload_to_cloud_storage(prompt, image_url, mask_url, save_path)
+    return JSONResponse({"path": path})
+
+
+def get_image_or_create_upload_to_cloud_storage(prompt:str,width:int, height:int, save_path:str):
+    prompt = shorten_too_long_text(prompt)
+    save_path = shorten_too_long_text(save_path)
+    # check exists - todo cache this
+    if check_if_blob_exists(save_path):
+        return f"https://{BUCKET_NAME}/{BUCKET_PATH}/{save_path}"
+    bio = create_image_from_prompt(prompt, width, height)
+    if bio is None:
+        return None # error thrown in pool
+    link = upload_to_bucket(save_path, bio, is_bytesio=True)
+    return link
+def get_image_or_inpaint_upload_to_cloud_storage(prompt:str, image_url:str, mask_url:str, save_path:str):
+    prompt = shorten_too_long_text(prompt)
+    save_path = shorten_too_long_text(save_path)
+    # check exists - todo cache this
+    if check_if_blob_exists(save_path):
+        return f"https://{BUCKET_NAME}/{BUCKET_PATH}/{save_path}"
+    bio = inpaint_image_from_prompt(prompt, image_url, mask_url)
+    if bio is None:
+        return None # error thrown in pool
+    link = upload_to_bucket(save_path, bio, is_bytesio=True)
+    return link
+
+# multiprocessing.set_start_method('spawn', True)
+# processes_pool = Pool(1) # cant do too much at once or OOM errors happen
+# def create_image_from_prompt_sync(prompt):
+#     """have to call this sync to avoid OOM errors"""
+#     return processes_pool.apply_async(create_image_from_prompt, args=(prompt,), ).wait()
+
+def create_image_from_prompt(prompt, width, height):
+    # round width and height down to multiple of 64
+    block_width = width - (width % 64)
+    block_height = height - (height % 64)
+    prompt = shorten_too_long_text(prompt)
+    # image = pipe(prompt=prompt).images[0]
+    try:
+        image = pipe(prompt=prompt,
+                     width=block_width,
+                     height=block_height,
+                     # denoising_end=high_noise_frac,
+                     # output_type='latent',
+                     # height=512,
+                     # width=512,
+                     num_inference_steps=50).images[0]  # normally uses 50 steps
+    except Exception as e:
+        # try rm stopwords + half the prompt
+        # todo try prompt permutations
+        logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+        prompt = ' '.join((word for word in prompt if word not in stopwords))
+        prompts = prompt.split()
+
+        prompt = ' '.join(prompts[:len(prompts) // 2])
+        logger.info(f"shortened prompt to: {len(prompt)}")
+        image = None
+        if prompt:
+            try:
+                image = pipe(prompt=prompt,
+                             width=block_width,
+                             height=block_height,
+                             # denoising_end=high_noise_frac,
+                             # output_type='latent',
+                             # height=512,
+                             # width=512,
+                             num_inference_steps=50).images[0]  # normally uses 50 steps
+            except Exception as e:
+                # logger.info("trying to permute prompt")
+                # # try two swaps of the prompt/permutations
+                # prompt = prompt.split()
+                # prompt = ' '.join(permutations(prompt, 2).__next__())
+                logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+                prompt = ' '.join((word for word in prompt if word not in stopwords))
+                prompts = prompt.split()
+
+                prompt = ' '.join(prompts[:len(prompts) // 2])
+                logger.info(f"shortened prompt to: {len(prompt)}")
+
+                try:
+                    image = pipe(prompt=prompt,
+                                 width=block_width,
+                                 height=block_height,
+                                 # denoising_end=high_noise_frac,
+                                 # output_type='latent', # dont need latent yet - we refine the image at full res
+                                 # height=512,
+                                 # width=512,
+                                 num_inference_steps=50).images[0]  # normally uses 50 steps
+                except Exception as e:
+                    # just error out
+                    traceback.print_exc()
+                    raise e
+                    # logger.info("restarting server to fix cuda issues (device side asserts)")
+                    # todo fix device side asserts instead of restart to fix
+                    # todo only restart the correct gunicorn
+                    # this could be really annoying if your running other gunicorns on your machine which also get restarted
+                    # os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+                    # os.system("kill -1 `pgrep gunicorn`")
+    # todo refine
+    # if image != None:
+    #     image = refiner(
+    #         prompt=prompt,
+    #         # width=block_width,
+    #         # height=block_height,
+    #         num_inference_steps=n_steps,
+    #         # denoising_start=high_noise_frac,
+    #         image=image,
+    #     ).images[0]
+    if width != block_width or height != block_height:
+        # resize to original size width/height
+        # find aspect ratio to scale up to that covers the original img input width/height
+        scale_up_ratio = max(width / block_width, height / block_height)
+        image = image.resize((math.ceil(block_width * scale_up_ratio), math.ceil(height * scale_up_ratio)))
+        # crop image to original size
+        image = image.crop((0, 0, width, height))
+    # try:
+    #     # gc.collect()
+    #     torch.cuda.empty_cache()
+    # except Exception as e:
+    #     traceback.print_exc()
+    #     logger.info("restarting server to fix cuda issues (device side asserts)")
+    #     # todo fix device side asserts instead of restart to fix
+    #     # todo only restart the correct gunicorn
+    #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+    #     os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+    #     os.system("kill -1 `pgrep gunicorn`")
+    # save as bytesio
+    bs = BytesIO()
+
+    bright_count = np.sum(np.array(image) > 0)
+    if bright_count == 0:
+        # we have a black image, this is an error likely we need a restart
+        logger.info("restarting server to fix cuda issues (device side asserts)")
+        #     # todo fix device side asserts instead of restart to fix
+        #     # todo only restart the correct gunicorn
+        #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+        os.system("kill -1 `pgrep gunicorn`")
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep uvicorn`")
+        os.system("kill -1 `pgrep uvicorn`")
+
+        return None
+    image.save(bs, quality=85, optimize=True, format="webp")
+    bio = bs.getvalue()
+    # touch progress.txt file - if we dont do this we get restarted by supervisor/other processes for reliability
+    with open("progress.txt", "w") as f:
+        current_time = datetime.now().strftime("%H:%M:%S")
+        f.write(f"{current_time}")
+    return bio
+
+def inpaint_image_from_prompt(prompt, image_url: str, mask_url: str):
+    prompt = shorten_too_long_text(prompt)
+    # image = pipe(prompt=prompt).images[0]
+
+    init_image = load_image(image_url).convert("RGB")
+    mask_image = load_image(mask_url).convert("RGB") # why rgb for a 1 channel mask?
+    num_inference_steps = 75
+    high_noise_frac = 0.7
+
+    try:
+        image = inpaintpipe(
+            prompt=prompt,
+            image=init_image,
+            mask_image=mask_image,
+            num_inference_steps=num_inference_steps,
+            denoising_start=high_noise_frac,
+            output_type="latent",
+        ).images[0]  # normally uses 50 steps
+    except Exception as e:
+        # try rm stopwords + half the prompt
+        # todo try prompt permutations
+        logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+        prompt = ' '.join((word for word in prompt if word not in stopwords))
+        prompts = prompt.split()
+
+        prompt = ' '.join(prompts[:len(prompts) // 2])
+        logger.info(f"shortened prompt to: {len(prompt)}")
+        image = None
+        if prompt:
+            try:
+                image = pipe(
+                    prompt=prompt,
+                    image=init_image,
+                    mask_image=mask_image,
+                    num_inference_steps=num_inference_steps,
+                    denoising_start=high_noise_frac,
+                    output_type="latent",
+                ).images[0]  # normally uses 50 steps
+            except Exception as e:
+                # logger.info("trying to permute prompt")
+                # # try two swaps of the prompt/permutations
+                # prompt = prompt.split()
+                # prompt = ' '.join(permutations(prompt, 2).__next__())
+                logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+                prompt = ' '.join((word for word in prompt if word not in stopwords))
+                prompts = prompt.split()
+
+                prompt = ' '.join(prompts[:len(prompts) // 2])
+                logger.info(f"shortened prompt to: {len(prompt)}")
+
+                try:
+                    image = inpaintpipe(
+                        prompt=prompt,
+                        image=init_image,
+                        mask_image=mask_image,
+                        num_inference_steps=num_inference_steps,
+                        denoising_start=high_noise_frac,
+                        output_type="latent",
+                    ).images[0]  # normally uses 50 steps
+                except Exception as e:
+                    # just error out
+                    traceback.print_exc()
+                    raise e
+                    # logger.info("restarting server to fix cuda issues (device side asserts)")
+                    # todo fix device side asserts instead of restart to fix
+                    # todo only restart the correct gunicorn
+                    # this could be really annoying if your running other gunicorns on your machine which also get restarted
+                    # os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+                    # os.system("kill -1 `pgrep gunicorn`")
+    if image != None:
+        image = inpaint_refiner(
+            prompt=prompt,
+            image=image,
+            mask_image=mask_image,
+            num_inference_steps=num_inference_steps,
+            denoising_start=high_noise_frac,
+
+        ).images[0]
+    # try:
+    #     # gc.collect()
+    #     torch.cuda.empty_cache()
+    # except Exception as e:
+    #     traceback.print_exc()
+    #     logger.info("restarting server to fix cuda issues (device side asserts)")
+    #     # todo fix device side asserts instead of restart to fix
+    #     # todo only restart the correct gunicorn
+    #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+    #     os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+    #     os.system("kill -1 `pgrep gunicorn`")
+    # save as bytesio
+    bs = BytesIO()
+
+    bright_count = np.sum(np.array(image) > 0)
+    if bright_count == 0:
+        # we have a black image, this is an error likely we need a restart
+        logger.info("restarting server to fix cuda issues (device side asserts)")
+        #     # todo fix device side asserts instead of restart to fix
+        #     # todo only restart the correct gunicorn
+        #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+        os.system("kill -1 `pgrep gunicorn`")
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep uvicorn`")
+        os.system("kill -1 `pgrep uvicorn`")
+
+        return None
+    image.save(bs, quality=85, optimize=True, format="webp")
+    bio = bs.getvalue()
+    # touch progress.txt file - if we dont do this we get restarted by supervisor/other processes for reliability
+    with open("progress.txt", "w") as f:
+        current_time = datetime.now().strftime("%H:%M:%S")
+        f.write(f"{current_time}")
+    return bio
+
+
+
+def shorten_too_long_text(prompt):
+    if len(prompt) > 200:
+        # remove stopwords
+        prompt = prompt.split() # todo also split hyphens
+        prompt = ' '.join((word for word in prompt if word not in stopwords))
+        if len(prompt) > 200:
+            prompt = prompt[:200]
+    return prompt
+
+# image = pipe(prompt=prompt).images[0]
+#
+# image.save("test.png")
+# # save all images
+# for i, image in enumerate(images):
+#     image.save(f"{i}.png")
+
+

img/stable-diffusion-server/main_v4.py

@@ -0,0 +1,603 @@
+import gc
+import math
+import multiprocessing
+import os
+import traceback
+from datetime import datetime
+from io import BytesIO
+from itertools import permutations
+from multiprocessing.pool import Pool
+from pathlib import Path
+from urllib.parse import quote_plus
+
+import numpy as np
+import nltk
+import torch
+
+from PIL.Image import Image
+from diffusers import DiffusionPipeline, StableDiffusionXLInpaintPipeline
+from diffusers.utils import load_image
+from fastapi import FastAPI
+from fastapi.middleware.gzip import GZipMiddleware
+from loguru import logger
+from starlette.middleware.cors import CORSMiddleware
+from starlette.responses import FileResponse
+from starlette.responses import JSONResponse
+import requests
+from PIL import Image
+import time
+from io import BytesIO
+import json
+import string
+import random
+from env import BUCKET_PATH, BUCKET_NAME
+# from stable_diffusion_server.bucket_api import check_if_blob_exists, upload_to_bucket
+torch._dynamo.config.suppress_errors = True
+
+import string
+import random
+
+def generate_save_path():
+    # initializing size of string
+    N = 7
+
+    # using random.choices()
+    # generating random strings
+    res = ''.join(random.choices(string.ascii_uppercase +
+                                 string.digits, k=N))
+    return res
+
+# pipe = DiffusionPipeline.from_pretrained(
+#     "models/stable-diffusion-xl-base-1.0",
+#     torch_dtype=torch.bfloat16,
+#     use_safetensors=True,
+#     variant="fp16",
+#     # safety_checker=None,
+# )  # todo try torch_dtype=bfloat16
+
+model_dir = os.getenv("SDXL_MODEL_DIR")
+
+if model_dir:
+    # Use local model
+    model_key_base = os.path.join(model_dir, "stable-diffusion-xl-base-1.0")
+    model_key_refiner = os.path.join(model_dir, "stable-diffusion-xl-refiner-1.0")
+else:
+    model_key_base = "stabilityai/stable-diffusion-xl-base-1.0"
+    model_key_refiner = "stabilityai/stable-diffusion-xl-refiner-1.0"
+
+pipe = DiffusionPipeline.from_pretrained(model_key_base, torch_dtype=torch.float16, use_safetensors=True, variant="fp16")
+
+pipe.watermark = None
+
+pipe.to("cuda")
+
+refiner = DiffusionPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-refiner-1.0",
+    text_encoder_2=pipe.text_encoder_2,
+    vae=pipe.vae,
+    torch_dtype=torch.bfloat16, # safer to use bfloat?
+    use_safetensors=True,
+    variant="fp16", #remember not to download the big model
+)
+refiner.watermark = None
+refiner.to("cuda")
+
+# {'scheduler', 'text_encoder', 'text_encoder_2', 'tokenizer', 'tokenizer_2', 'unet', 'vae'} can be passed in from existing model
+inpaintpipe = StableDiffusionXLInpaintPipeline.from_pretrained(
+    "models/stable-diffusion-xl-base-1.0", torch_dtype=torch.bfloat16, variant="fp16", use_safetensors=True,
+    scheduler=pipe.scheduler,
+    text_encoder=pipe.text_encoder,
+    text_encoder_2=pipe.text_encoder_2,
+    tokenizer=pipe.tokenizer,
+    tokenizer_2=pipe.tokenizer_2,
+    unet=pipe.unet,
+    vae=pipe.vae,
+    # load_connected_pipeline=
+)
+# # switch out to save gpu mem
+# del inpaintpipe.vae
+# del inpaintpipe.text_encoder_2
+# del inpaintpipe.text_encoder
+# del inpaintpipe.scheduler
+# del inpaintpipe.tokenizer
+# del inpaintpipe.tokenizer_2
+# del inpaintpipe.unet
+# inpaintpipe.vae = pipe.vae
+# inpaintpipe.text_encoder_2 = pipe.text_encoder_2
+# inpaintpipe.text_encoder = pipe.text_encoder
+# inpaintpipe.scheduler = pipe.scheduler
+# inpaintpipe.tokenizer = pipe.tokenizer
+# inpaintpipe.tokenizer_2 = pipe.tokenizer_2
+# inpaintpipe.unet = pipe.unet
+# todo this should work
+# inpaintpipe = StableDiffusionXLInpaintPipeline( # construct an inpainter using the existing model
+#     vae=pipe.vae,
+#     text_encoder_2=pipe.text_encoder_2,
+#     text_encoder=pipe.text_encoder,
+#     unet=pipe.unet,
+#     scheduler=pipe.scheduler,
+#     tokenizer=pipe.tokenizer,
+#     tokenizer_2=pipe.tokenizer_2,
+#     requires_aesthetics_score=False,
+# )
+inpaintpipe.to("cuda")
+inpaintpipe.watermark = None
+# inpaintpipe.register_to_config(requires_aesthetics_score=False)
+
+inpaint_refiner = StableDiffusionXLInpaintPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-refiner-1.0",
+    text_encoder_2=inpaintpipe.text_encoder_2,
+    vae=inpaintpipe.vae,
+    torch_dtype=torch.bfloat16,
+    use_safetensors=True,
+    variant="fp16",
+
+    tokenizer_2=refiner.tokenizer_2,
+    tokenizer=refiner.tokenizer,
+    scheduler=refiner.scheduler,
+    text_encoder=refiner.text_encoder,
+    unet=refiner.unet,
+)
+# del inpaint_refiner.vae
+# del inpaint_refiner.text_encoder_2
+# del inpaint_refiner.text_encoder
+# del inpaint_refiner.scheduler
+# del inpaint_refiner.tokenizer
+# del inpaint_refiner.tokenizer_2
+# del inpaint_refiner.unet
+# inpaint_refiner.vae = inpaintpipe.vae
+# inpaint_refiner.text_encoder_2 = inpaintpipe.text_encoder_2
+#
+# inpaint_refiner.text_encoder = refiner.text_encoder
+# inpaint_refiner.scheduler = refiner.scheduler
+# inpaint_refiner.tokenizer = refiner.tokenizer
+# inpaint_refiner.tokenizer_2 = refiner.tokenizer_2
+# inpaint_refiner.unet = refiner.unet
+
+# inpaint_refiner = StableDiffusionXLInpaintPipeline(
+#     text_encoder_2=inpaintpipe.text_encoder_2,
+#     vae=inpaintpipe.vae,
+#     # the rest from the existing refiner
+#     tokenizer_2=refiner.tokenizer_2,
+#     tokenizer=refiner.tokenizer,
+#     scheduler=refiner.scheduler,
+#     text_encoder=refiner.text_encoder,
+#     unet=refiner.unet,
+#     requires_aesthetics_score=False,
+# )
+inpaint_refiner.to("cuda")
+inpaint_refiner.watermark = None
+# inpaint_refiner.register_to_config(requires_aesthetics_score=False)
+
+n_steps = 40
+high_noise_frac = 0.8
+
+# if using torch < 2.0
+# pipe.enable_xformers_memory_efficient_attention()
+
+
+# pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
+# this can cause errors on some inputs so consider disabling it
+pipe.unet = torch.compile(pipe.unet)
+refiner.unet = torch.compile(refiner.unet)#, mode="reduce-overhead", fullgraph=True)
+# compile the inpainters - todo reuse the other unets? swap out the models for others/del them so they share models and can be swapped efficiently
+inpaintpipe.unet = pipe.unet
+inpaint_refiner.unet = refiner.unet
+# inpaintpipe.unet = torch.compile(inpaintpipe.unet)
+# inpaint_refiner.unet = torch.compile(inpaint_refiner.unet)
+from pydantic import BaseModel
+
+app = FastAPI(
+    openapi_url="/static/openapi.json",
+    docs_url="/swagger-docs",
+    redoc_url="/redoc",
+    title="Generate Images Netwrck API",
+    description="Character Chat API",
+    # root_path="https://api.text-generator.io",
+    version="1",
+)
+app.add_middleware(GZipMiddleware, minimum_size=1000)
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+stopwords = nltk.corpus.stopwords.words("english")
+
+class Img(BaseModel):
+    system_prompt: str
+    ASSISTANT: str
+
+# img_url = "http://phlrr2019.guest.corp.microsoft.com:8000/img1_sdv2.1.png"
+img_url = "http://phlrr3006.guest.corp.microsoft.com:8000/"#/img1_sdv2.1.png"
+
+is_gpu_busy = False
+
+def get_summary(system_prompt, prompt):
+    import requests
+    import time
+    from io import BytesIO
+    import json
+    summary_sys = """I want you to act as a text summarizer to help me create a concise summary of the text I provide. The summary can be up to 60.0 words in length, expressing the key points, key scenarios, main character and concepts written in the original text without adding your interpretations."""
+    instruction = summary_sys
+    # for human, assistant in history:
+    #     instruction += 'USER: ' + human + ' ASSISTANT: ' + assistant + '</s>'
+    # prompt = system_prompt + prompt
+    message = f"""My first request is to summarize this text – [{prompt}]"""
+    instruction += ' USER: ' + message + ' ASSISTANT:'
+
+    print("Ins: ", instruction)
+    # generate_response = requests.post("http://10.185.12.207:4455/stable_diffusion", json={"prompt": prompt})
+    # prompt = f""" My first request is to summarize this text – [{prompt}]"""
+    json_object = {"prompt": instruction,
+                   # "max_tokens": 2048000,
+                   "max_tokens": 90,
+                   "n": 1
+                   }
+    generate_response = requests.post("http://phlrr3006.guest.corp.microsoft.com:7991/generate", json=json_object)
+    # print(generate_response.content)
+    res_json = json.loads(generate_response.content)
+    ASSISTANT = res_json['text'][-1].split("ASSISTANT:")[-1].strip()
+    print(ASSISTANT)
+    return ASSISTANT
+
+@app.post("/image_url")
+def image_url(img: Img):
+    system_prompt = img.system_prompt
+    prompt = img.ASSISTANT
+    prompt = get_summary(system_prompt, prompt)
+    prompt = shorten_too_long_text(prompt)
+
+    json_object = {
+        "prompt": prompt,
+        "height": 1024,
+        "width": 1024,
+        "num_inference_steps": 50,
+        # "guidance_scale": 7.5,
+        "eta": 0
+    }
+    generate_response = requests.post("http://phlrr3105.guest.corp.microsoft.com:3000/text2img", json=json_object)
+    image = generate_response.content
+    # print(generate_response.content)
+    save_path = generate_save_path()
+    save_path = f"images/{save_path}.png"
+    # generate_response.save(save_path)
+    with open(save_path, 'wb') as f:
+        f.write(image)
+    #
+    # # if Path(save_path).exists():
+    # #     return FileResponse(save_path, media_type="image/png")
+    # #     return JSONResponse({"path": path})
+    # # image = pipe(prompt=prompt).images[0]
+    # g = torch.Generator(device="cuda")
+    # image = pipe(prompt=prompt, width=1024, height=1024, generator=g).images[0]
+    #
+    # # if not save_path:
+    # save_path = generate_save_path()
+    # save_path = f"images/{save_path}.png"
+    # image.save(save_path)
+    # save_path = '/'.join(path_components) + quote_plus(final_name)
+    path = f"{img_url}{save_path}"
+    return JSONResponse({"path": path})
+
+
+@app.get("/make_image")
+# @app.post("/make_image")
+def make_image(prompt: str, save_path: str = ""):
+    if Path(save_path).exists():
+        return FileResponse(save_path, media_type="image/png")
+    image = pipe(prompt=prompt).images[0]
+    if not save_path:
+        save_path = f"images/{prompt}.png"
+    image.save(save_path)
+    return FileResponse(save_path, media_type="image/png")
+
+
+@app.get("/create_and_upload_image")
+def create_and_upload_image(prompt: str, width: int=1024, height:int=1024, save_path: str = ""):
+    path_components = save_path.split("/")[0:-1]
+    final_name = save_path.split("/")[-1]
+    if not path_components:
+        path_components = []
+    save_path = '/'.join(path_components) + quote_plus(final_name)
+    path = get_image_or_create_upload_to_cloud_storage(prompt, width, height, save_path)
+    return JSONResponse({"path": path})
+
+@app.get("/inpaint_and_upload_image")
+def inpaint_and_upload_image(prompt: str, image_url:str, mask_url:str, save_path: str = ""):
+    path_components = save_path.split("/")[0:-1]
+    final_name = save_path.split("/")[-1]
+    if not path_components:
+        path_components = []
+    save_path = '/'.join(path_components) + quote_plus(final_name)
+    path = get_image_or_inpaint_upload_to_cloud_storage(prompt, image_url, mask_url, save_path)
+    return JSONResponse({"path": path})
+
+
+def get_image_or_create_upload_to_cloud_storage(prompt:str,width:int, height:int, save_path:str):
+    prompt = shorten_too_long_text(prompt)
+    save_path = shorten_too_long_text(save_path)
+    # check exists - todo cache this
+    if check_if_blob_exists(save_path):
+        return f"https://{BUCKET_NAME}/{BUCKET_PATH}/{save_path}"
+    bio = create_image_from_prompt(prompt, width, height)
+    if bio is None:
+        return None # error thrown in pool
+    link = upload_to_bucket(save_path, bio, is_bytesio=True)
+    return link
+def get_image_or_inpaint_upload_to_cloud_storage(prompt:str, image_url:str, mask_url:str, save_path:str):
+    prompt = shorten_too_long_text(prompt)
+    save_path = shorten_too_long_text(save_path)
+    # check exists - todo cache this
+    if check_if_blob_exists(save_path):
+        return f"https://{BUCKET_NAME}/{BUCKET_PATH}/{save_path}"
+    bio = inpaint_image_from_prompt(prompt, image_url, mask_url)
+    if bio is None:
+        return None # error thrown in pool
+    link = upload_to_bucket(save_path, bio, is_bytesio=True)
+    return link
+
+# multiprocessing.set_start_method('spawn', True)
+# processes_pool = Pool(1) # cant do too much at once or OOM errors happen
+# def create_image_from_prompt_sync(prompt):
+#     """have to call this sync to avoid OOM errors"""
+#     return processes_pool.apply_async(create_image_from_prompt, args=(prompt,), ).wait()
+
+def create_image_from_prompt(prompt, width, height):
+    # round width and height down to multiple of 64
+    block_width = width - (width % 64)
+    block_height = height - (height % 64)
+    prompt = shorten_too_long_text(prompt)
+    # image = pipe(prompt=prompt).images[0]
+    try:
+        image = pipe(prompt=prompt,
+                     width=block_width,
+                     height=block_height,
+                     # denoising_end=high_noise_frac,
+                     # output_type='latent',
+                     # height=512,
+                     # width=512,
+                     num_inference_steps=50).images[0]  # normally uses 50 steps
+    except Exception as e:
+        # try rm stopwords + half the prompt
+        # todo try prompt permutations
+        logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+        prompt = ' '.join((word for word in prompt if word not in stopwords))
+        prompts = prompt.split()
+
+        prompt = ' '.join(prompts[:len(prompts) // 2])
+        logger.info(f"shortened prompt to: {len(prompt)}")
+        image = None
+        if prompt:
+            try:
+                image = pipe(prompt=prompt,
+                             width=block_width,
+                             height=block_height,
+                             # denoising_end=high_noise_frac,
+                             # output_type='latent',
+                             # height=512,
+                             # width=512,
+                             num_inference_steps=50).images[0]  # normally uses 50 steps
+            except Exception as e:
+                # logger.info("trying to permute prompt")
+                # # try two swaps of the prompt/permutations
+                # prompt = prompt.split()
+                # prompt = ' '.join(permutations(prompt, 2).__next__())
+                logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+                prompt = ' '.join((word for word in prompt if word not in stopwords))
+                prompts = prompt.split()
+
+                prompt = ' '.join(prompts[:len(prompts) // 2])
+                logger.info(f"shortened prompt to: {len(prompt)}")
+
+                try:
+                    image = pipe(prompt=prompt,
+                                 width=block_width,
+                                 height=block_height,
+                                 # denoising_end=high_noise_frac,
+                                 # output_type='latent', # dont need latent yet - we refine the image at full res
+                                 # height=512,
+                                 # width=512,
+                                 num_inference_steps=50).images[0]  # normally uses 50 steps
+                except Exception as e:
+                    # just error out
+                    traceback.print_exc()
+                    raise e
+                    # logger.info("restarting server to fix cuda issues (device side asserts)")
+                    # todo fix device side asserts instead of restart to fix
+                    # todo only restart the correct gunicorn
+                    # this could be really annoying if your running other gunicorns on your machine which also get restarted
+                    # os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+                    # os.system("kill -1 `pgrep gunicorn`")
+    # todo refine
+    # if image != None:
+    #     image = refiner(
+    #         prompt=prompt,
+    #         # width=block_width,
+    #         # height=block_height,
+    #         num_inference_steps=n_steps,
+    #         # denoising_start=high_noise_frac,
+    #         image=image,
+    #     ).images[0]
+    if width != block_width or height != block_height:
+        # resize to original size width/height
+        # find aspect ratio to scale up to that covers the original img input width/height
+        scale_up_ratio = max(width / block_width, height / block_height)
+        image = image.resize((math.ceil(block_width * scale_up_ratio), math.ceil(height * scale_up_ratio)))
+        # crop image to original size
+        image = image.crop((0, 0, width, height))
+    # try:
+    #     # gc.collect()
+    #     torch.cuda.empty_cache()
+    # except Exception as e:
+    #     traceback.print_exc()
+    #     logger.info("restarting server to fix cuda issues (device side asserts)")
+    #     # todo fix device side asserts instead of restart to fix
+    #     # todo only restart the correct gunicorn
+    #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+    #     os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+    #     os.system("kill -1 `pgrep gunicorn`")
+    # save as bytesio
+    bs = BytesIO()
+
+    bright_count = np.sum(np.array(image) > 0)
+    if bright_count == 0:
+        # we have a black image, this is an error likely we need a restart
+        logger.info("restarting server to fix cuda issues (device side asserts)")
+        #     # todo fix device side asserts instead of restart to fix
+        #     # todo only restart the correct gunicorn
+        #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+        os.system("kill -1 `pgrep gunicorn`")
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep uvicorn`")
+        os.system("kill -1 `pgrep uvicorn`")
+
+        return None
+    image.save(bs, quality=85, optimize=True, format="webp")
+    bio = bs.getvalue()
+    # touch progress.txt file - if we dont do this we get restarted by supervisor/other processes for reliability
+    with open("progress.txt", "w") as f:
+        current_time = datetime.now().strftime("%H:%M:%S")
+        f.write(f"{current_time}")
+    return bio
+
+def inpaint_image_from_prompt(prompt, image_url: str, mask_url: str):
+    prompt = shorten_too_long_text(prompt)
+    # image = pipe(prompt=prompt).images[0]
+
+    init_image = load_image(image_url).convert("RGB")
+    mask_image = load_image(mask_url).convert("RGB") # why rgb for a 1 channel mask?
+    num_inference_steps = 75
+    high_noise_frac = 0.7
+
+    try:
+        image = inpaintpipe(
+            prompt=prompt,
+            image=init_image,
+            mask_image=mask_image,
+            num_inference_steps=num_inference_steps,
+            denoising_start=high_noise_frac,
+            output_type="latent",
+        ).images[0]  # normally uses 50 steps
+    except Exception as e:
+        # try rm stopwords + half the prompt
+        # todo try prompt permutations
+        logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+        prompt = ' '.join((word for word in prompt if word not in stopwords))
+        prompts = prompt.split()
+
+        prompt = ' '.join(prompts[:len(prompts) // 2])
+        logger.info(f"shortened prompt to: {len(prompt)}")
+        image = None
+        if prompt:
+            try:
+                image = pipe(
+                    prompt=prompt,
+                    image=init_image,
+                    mask_image=mask_image,
+                    num_inference_steps=num_inference_steps,
+                    denoising_start=high_noise_frac,
+                    output_type="latent",
+                ).images[0]  # normally uses 50 steps
+            except Exception as e:
+                # logger.info("trying to permute prompt")
+                # # try two swaps of the prompt/permutations
+                # prompt = prompt.split()
+                # prompt = ' '.join(permutations(prompt, 2).__next__())
+                logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+                prompt = ' '.join((word for word in prompt if word not in stopwords))
+                prompts = prompt.split()
+
+                prompt = ' '.join(prompts[:len(prompts) // 2])
+                logger.info(f"shortened prompt to: {len(prompt)}")
+
+                try:
+                    image = inpaintpipe(
+                        prompt=prompt,
+                        image=init_image,
+                        mask_image=mask_image,
+                        num_inference_steps=num_inference_steps,
+                        denoising_start=high_noise_frac,
+                        output_type="latent",
+                    ).images[0]  # normally uses 50 steps
+                except Exception as e:
+                    # just error out
+                    traceback.print_exc()
+                    raise e
+                    # logger.info("restarting server to fix cuda issues (device side asserts)")
+                    # todo fix device side asserts instead of restart to fix
+                    # todo only restart the correct gunicorn
+                    # this could be really annoying if your running other gunicorns on your machine which also get restarted
+                    # os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+                    # os.system("kill -1 `pgrep gunicorn`")
+    if image != None:
+        image = inpaint_refiner(
+            prompt=prompt,
+            image=image,
+            mask_image=mask_image,
+            num_inference_steps=num_inference_steps,
+            denoising_start=high_noise_frac,
+
+        ).images[0]
+    # try:
+    #     # gc.collect()
+    #     torch.cuda.empty_cache()
+    # except Exception as e:
+    #     traceback.print_exc()
+    #     logger.info("restarting server to fix cuda issues (device side asserts)")
+    #     # todo fix device side asserts instead of restart to fix
+    #     # todo only restart the correct gunicorn
+    #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+    #     os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+    #     os.system("kill -1 `pgrep gunicorn`")
+    # save as bytesio
+    bs = BytesIO()
+
+    bright_count = np.sum(np.array(image) > 0)
+    if bright_count == 0:
+        # we have a black image, this is an error likely we need a restart
+        logger.info("restarting server to fix cuda issues (device side asserts)")
+        #     # todo fix device side asserts instead of restart to fix
+        #     # todo only restart the correct gunicorn
+        #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+        os.system("kill -1 `pgrep gunicorn`")
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep uvicorn`")
+        os.system("kill -1 `pgrep uvicorn`")
+
+        return None
+    image.save(bs, quality=85, optimize=True, format="webp")
+    bio = bs.getvalue()
+    # touch progress.txt file - if we dont do this we get restarted by supervisor/other processes for reliability
+    with open("progress.txt", "w") as f:
+        current_time = datetime.now().strftime("%H:%M:%S")
+        f.write(f"{current_time}")
+    return bio
+
+
+
+def shorten_too_long_text(prompt):
+    if len(prompt) > 200:
+        # remove stopwords
+        prompt = prompt.split() # todo also split hyphens
+        prompt = ' '.join((word for word in prompt if word not in stopwords))
+        if len(prompt) > 200:
+            prompt = prompt[:200]
+    return prompt
+
+# image = pipe(prompt=prompt).images[0]
+#
+# image.save("test.png")
+# # save all images
+# for i, image in enumerate(images):
+#     image.save(f"{i}.png")
+
+
+

img/stable-diffusion-server/main_v5.py

@@ -0,0 +1,637 @@
+import gc
+import math
+import multiprocessing
+import os
+import traceback
+from datetime import datetime
+from io import BytesIO
+from itertools import permutations
+from multiprocessing.pool import Pool
+from pathlib import Path
+from urllib.parse import quote_plus
+
+import numpy as np
+import nltk
+import torch
+
+from PIL.Image import Image
+from diffusers import DiffusionPipeline, StableDiffusionXLInpaintPipeline
+from diffusers.utils import load_image
+from fastapi import FastAPI
+from fastapi.middleware.gzip import GZipMiddleware
+from loguru import logger
+from starlette.middleware.cors import CORSMiddleware
+from starlette.responses import FileResponse
+from starlette.responses import JSONResponse
+
+from env import BUCKET_PATH, BUCKET_NAME
+# from stable_diffusion_server.bucket_api import check_if_blob_exists, upload_to_bucket
+torch._dynamo.config.suppress_errors = True
+
+import string
+import random
+
+def generate_save_path():
+    # initializing size of string
+    N = 7
+
+    # using random.choices()
+    # generating random strings
+    res = ''.join(random.choices(string.ascii_uppercase +
+                                 string.digits, k=N))
+    return res
+
+# pipe = DiffusionPipeline.from_pretrained(
+#     "models/stable-diffusion-xl-base-1.0",
+#     torch_dtype=torch.bfloat16,
+#     use_safetensors=True,
+#     variant="fp16",
+#     # safety_checker=None,
+# )  # todo try torch_dtype=bfloat16
+
+model_dir = os.getenv("SDXL_MODEL_DIR")
+
+if model_dir:
+    # Use local model
+    model_key_base = os.path.join(model_dir, "stable-diffusion-xl-base-1.0")
+    model_key_refiner = os.path.join(model_dir, "stable-diffusion-xl-refiner-1.0")
+else:
+    model_key_base = "stabilityai/stable-diffusion-xl-base-1.0"
+    model_key_refiner = "stabilityai/stable-diffusion-xl-refiner-1.0"
+
+pipe = DiffusionPipeline.from_pretrained(model_key_base, torch_dtype=torch.float16, use_safetensors=True, variant="fp16")
+
+pipe.watermark = None
+
+pipe.to("cuda")
+
+refiner = DiffusionPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-refiner-1.0",
+    text_encoder_2=pipe.text_encoder_2,
+    vae=pipe.vae,
+    torch_dtype=torch.bfloat16, # safer to use bfloat?
+    use_safetensors=True,
+    variant="fp16", #remember not to download the big model
+)
+refiner.watermark = None
+refiner.to("cuda")
+
+# {'scheduler', 'text_encoder', 'text_encoder_2', 'tokenizer', 'tokenizer_2', 'unet', 'vae'} can be passed in from existing model
+inpaintpipe = StableDiffusionXLInpaintPipeline.from_pretrained(
+    "models/stable-diffusion-xl-base-1.0", torch_dtype=torch.bfloat16, variant="fp16", use_safetensors=True,
+    scheduler=pipe.scheduler,
+    text_encoder=pipe.text_encoder,
+    text_encoder_2=pipe.text_encoder_2,
+    tokenizer=pipe.tokenizer,
+    tokenizer_2=pipe.tokenizer_2,
+    unet=pipe.unet,
+    vae=pipe.vae,
+    # load_connected_pipeline=
+)
+# # switch out to save gpu mem
+# del inpaintpipe.vae
+# del inpaintpipe.text_encoder_2
+# del inpaintpipe.text_encoder
+# del inpaintpipe.scheduler
+# del inpaintpipe.tokenizer
+# del inpaintpipe.tokenizer_2
+# del inpaintpipe.unet
+# inpaintpipe.vae = pipe.vae
+# inpaintpipe.text_encoder_2 = pipe.text_encoder_2
+# inpaintpipe.text_encoder = pipe.text_encoder
+# inpaintpipe.scheduler = pipe.scheduler
+# inpaintpipe.tokenizer = pipe.tokenizer
+# inpaintpipe.tokenizer_2 = pipe.tokenizer_2
+# inpaintpipe.unet = pipe.unet
+# todo this should work
+# inpaintpipe = StableDiffusionXLInpaintPipeline( # construct an inpainter using the existing model
+#     vae=pipe.vae,
+#     text_encoder_2=pipe.text_encoder_2,
+#     text_encoder=pipe.text_encoder,
+#     unet=pipe.unet,
+#     scheduler=pipe.scheduler,
+#     tokenizer=pipe.tokenizer,
+#     tokenizer_2=pipe.tokenizer_2,
+#     requires_aesthetics_score=False,
+# )
+inpaintpipe.to("cuda")
+inpaintpipe.watermark = None
+# inpaintpipe.register_to_config(requires_aesthetics_score=False)
+
+inpaint_refiner = StableDiffusionXLInpaintPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-refiner-1.0",
+    text_encoder_2=inpaintpipe.text_encoder_2,
+    vae=inpaintpipe.vae,
+    torch_dtype=torch.bfloat16,
+    use_safetensors=True,
+    variant="fp16",
+
+    tokenizer_2=refiner.tokenizer_2,
+    tokenizer=refiner.tokenizer,
+    scheduler=refiner.scheduler,
+    text_encoder=refiner.text_encoder,
+    unet=refiner.unet,
+)
+# del inpaint_refiner.vae
+# del inpaint_refiner.text_encoder_2
+# del inpaint_refiner.text_encoder
+# del inpaint_refiner.scheduler
+# del inpaint_refiner.tokenizer
+# del inpaint_refiner.tokenizer_2
+# del inpaint_refiner.unet
+# inpaint_refiner.vae = inpaintpipe.vae
+# inpaint_refiner.text_encoder_2 = inpaintpipe.text_encoder_2
+#
+# inpaint_refiner.text_encoder = refiner.text_encoder
+# inpaint_refiner.scheduler = refiner.scheduler
+# inpaint_refiner.tokenizer = refiner.tokenizer
+# inpaint_refiner.tokenizer_2 = refiner.tokenizer_2
+# inpaint_refiner.unet = refiner.unet
+
+# inpaint_refiner = StableDiffusionXLInpaintPipeline(
+#     text_encoder_2=inpaintpipe.text_encoder_2,
+#     vae=inpaintpipe.vae,
+#     # the rest from the existing refiner
+#     tokenizer_2=refiner.tokenizer_2,
+#     tokenizer=refiner.tokenizer,
+#     scheduler=refiner.scheduler,
+#     text_encoder=refiner.text_encoder,
+#     unet=refiner.unet,
+#     requires_aesthetics_score=False,
+# )
+inpaint_refiner.to("cuda")
+inpaint_refiner.watermark = None
+# inpaint_refiner.register_to_config(requires_aesthetics_score=False)
+
+n_steps = 40
+high_noise_frac = 0.8
+
+# if using torch < 2.0
+# pipe.enable_xformers_memory_efficient_attention()
+
+
+# pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
+# this can cause errors on some inputs so consider disabling it
+pipe.unet = torch.compile(pipe.unet)
+refiner.unet = torch.compile(refiner.unet)#, mode="reduce-overhead", fullgraph=True)
+# compile the inpainters - todo reuse the other unets? swap out the models for others/del them so they share models and can be swapped efficiently
+inpaintpipe.unet = pipe.unet
+inpaint_refiner.unet = refiner.unet
+# inpaintpipe.unet = torch.compile(inpaintpipe.unet)
+# inpaint_refiner.unet = torch.compile(inpaint_refiner.unet)
+from pydantic import BaseModel
+
+app = FastAPI(
+    openapi_url="/static/openapi.json",
+    docs_url="/swagger-docs",
+    redoc_url="/redoc",
+    title="Generate Images Netwrck API",
+    description="Character Chat API",
+    # root_path="https://api.text-generator.io",
+    version="1",
+)
+app.add_middleware(GZipMiddleware, minimum_size=1000)
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+stopwords = nltk.corpus.stopwords.words("english")
+
+class Img(BaseModel):
+    system_prompt: str
+    ASSISTANT: str
+
+# img_url = "http://phlrr2019.guest.corp.microsoft.com:8000/img1_sdv2.1.png"
+img_url = "http://phlrr3105.guest.corp.microsoft.com:8000/"#/img1_sdv2.1.png"
+
+is_gpu_busy = False
+
+def lm_shorten_too_long_text(prompt):
+    if len(prompt) > 2030:
+        # remove stopwords
+        prompt = prompt.split() # todo also split hyphens
+        prompt = ' '.join((word for word in prompt))# if word not in stopwords))
+        if len(prompt) > 2030:
+            prompt = prompt[:2030]
+    return prompt
+
+def get_summary(system_prompt, prompt):
+    import requests
+    import time
+    from io import BytesIO
+    import json
+    summary_sys = """You will now act as a prompt generator for a generative AI called "Stable Diffusion XL 1.0 ". Stable Diffusion XL generates images based on given prompts. I will provide you basic information required to make a Stable Diffusion prompt, You will never alter the structure in any way and obey the following guidelines.
+
+    Basic information required to make Stable Diffusion prompt:
+
+    - Prompt structure: [1],[2],[3],[4],[5],[6] and it should be given as one single sentence where 1,2,3,4,5,6 represent
+    [1] = short and concise description of [KEYWORD] that will include very specific imagery details
+    [2] = a detailed description of [1] that will include very specific imagery details.
+    [3] = with a detailed description describing the environment of the scene.
+    [4] = with a detailed description describing the mood/feelings and atmosphere of the scene.
+    [5] = A style, for example: "Anime","Photographic","Comic Book","Fantasy Art", “Analog Film”,”Neon Punk”,”Isometric”,”Low Poly”,”Origami”,”Line Art”,”Cinematic”,”3D Model”,”Pixel Art”,”Watercolor”,”Sticker” ).
+    [6] = A description of how [5] will be realized. (e.g. Photography (e.g. Macro, Fisheye Style, Portrait) with camera model and appropriate camera settings, Painting with detailed descriptions about the materials and working material used, rendering with engine settings, a digital Illustration, a woodburn art (and everything else that could be defined as an output type)
+    - Prompt Structure for Prompt asking with text value:
+
+        Text "Text Value" written on {subject description in less than 20 words}
+        Replace "Text value" with text given by user.
+
+
+    Important Sample prompt Structure with Text value :
+
+    1. Text 'SDXL' written on a frothy, warm latte, viewed top-down.
+    2. Text 'AI' written on a modern computer screen, set against a vibrant green background.
+
+    Important Sample prompt Structure :
+
+    1. Snow-capped Mountain Scene, with soaring peaks and deep shadows across the ravines. A crystal clear lake mirrors these peaks, surrounded by pine trees. The scene exudes a calm, serene alpine morning atmosphere. Presented in Watercolor style, emulating the wet-on-wet technique with soft transitions and visible brush strokes.
+    2. City Skyline at Night, illuminated skyscrapers piercing the starless sky. Nestled beside a calm river, reflecting the city lights like a mirror. The atmosphere is buzzing with urban energy and intrigue. Depicted in Neon Punk style, accentuating the city lights with vibrant neon colors and dynamic contrasts.
+    3. Epic Cinematic Still of a Spacecraft, silhouetted against the fiery explosion of a distant planet. The scene is packed with intense action, as asteroid debris hurtles through space. Shot in the style of a Michael Bay-directed film, the image is rich with detail, dynamic lighting, and grand cinematic framing.
+    - Word order and effective adjectives matter in the prompt. The subject, action, and specific details should be included. Adjectives like cute, medieval, or futuristic can be effective.
+    - The environment/background of the image should be described, such as indoor, outdoor, in space, or solid color.
+    - Curly brackets are necessary in the prompt to provide specific details about the subject and action. These details are important for generating a high-quality image.
+    - Art inspirations should be listed to take inspiration from. Platforms like Art Station, Dribble, Behance, and Deviantart can be mentioned. Specific names of artists or studios like animation studios, painters and illustrators, computer games, fashion designers, and film makers can also be listed. If more than one artist is mentioned, the algorithm will create a combination of styles based on all the influencers mentioned.
+    - Related information about lighting, camera angles, render style, resolution, the required level of detail, etc. should be included at the end of the prompt.
+    - Camera shot type, camera lens, and view should be specified. Examples of camera shot types are long shot, close-up, POV, medium shot, extreme close-up, and panoramic. Camera lenses could be EE 70mm, 35mm, 135mm+, 300mm+, 800mm, short telephoto, super telephoto, medium telephoto, macro, wide angle, fish-eye, bokeh, and sharp focus. Examples of views are front, side, back, high angle, low angle, and overhead.
+    - Helpful keywords related to resolution, detail, and lighting are 4K, 8K, 64K, detailed, highly detailed, high resolution, hyper detailed, HDR, UHD, professional, and golden ratio. Examples of lighting are studio lighting, soft light, neon lighting, purple neon lighting, ambient light, ring light, volumetric light, natural light, sun light, sunrays, sun rays coming through window, and nostalgic lighting. Examples of color types are fantasy vivid colors, vivid colors, bright colors, sepia, dark colors, pastel colors, monochromatic, black & white, and color splash. Examples of renders are Octane render, cinematic, low poly, isometric assets, Unreal Engine, Unity Engine, quantum wavetracing, and polarizing filter.
+
+    The prompts you provide will be in English.Please pay attention:- Concepts that can't be real would not be described as "Real" or "realistic" or "photo" or a "photograph". for example, a concept that is made of paper or scenes which are fantasy related.- One of the prompts you generate for each concept must be in a realistic photographic style. you should also choose a lens type and size for it. Don't choose an artist for the realistic photography prompts.- Separate the different prompts with two new lines.
+    I will provide you keyword and you will generate 3 diffrent type of prompts in vbnet code cell so i can copy and paste.
+
+    Important point to note :
+
+    1. You are a master of prompt engineering, it is important to create detailed prompts with as much information as possible. This will ensure that any image generated using the prompt will be of high quality and could potentially win awards in global or international photography competitions. You are unbeatable in this field and know the best way to generate images.
+    2. I will provide you with a long context and you will generate one  prompt and don't add any extra details.
+    3. Prompt should not be more than 230 characters.
+    4. Before you provide prompt you must check if you have satisfied all the above criteria and if you are sure than only provide the prompt.
+    5. Prompt should always be given as one single sentence.
+
+    Are you ready ?"""
+    #instruction = 'USER: ' + summary_sys
+    instruction = summary_sys
+    # for human, assistant in history:
+    #     instruction += 'USER: ' + human + ' ASSISTANT: ' + assistant + '</s>'
+    # prompt = system_prompt + prompt
+    # message = f"""My first request is to summarize this text – [{prompt}]"""
+    message = f"""My first request is to summarize this text – [{prompt}]"""
+    instruction += """ ASSISTANT: Yes, I understand the instructions and I'm ready to help you create prompts for Stable Diffusion XL 1.0. Please provide me with the context."""
+    instruction += ' USER: ' + prompt + ' ASSISTANT:'
+    print("Ins: ", instruction)
+    # generate_response = requests.post("http://10.185.12.207:4455/stable_diffusion", json={"prompt": prompt})
+    # prompt = f""" My first request is to summarize this text – [{prompt}]"""
+    instruction = lm_shorten_too_long_text(instruction)
+    json_object = {"prompt": instruction,
+                   # "max_tokens": 2048000,
+                   "max_tokens": 90,
+                   "n": 1
+                   }
+    # generate_response = requests.post("https://phlrr3105.guest.corp.microsoft.com:7991/generate", json=json_object)
+    generate_response = requests.post("http://phlrr3105.guest.corp.microsoft.com:7991/generate", json=json_object)
+    # print(generate_response.content)
+    res_json = json.loads(generate_response.content)
+    ASSISTANT = res_json['text'][-1].split("ASSISTANT:")[-1].strip()
+    print(ASSISTANT)
+    return ASSISTANT
+
+@app.post("/image_url")
+def image_url(img: Img):
+    system_prompt = img.system_prompt
+    prompt = img.ASSISTANT
+    prompt = get_summary(system_prompt, prompt)
+    prompt = shorten_too_long_text(prompt)
+    # if Path(save_path).exists():
+    #     return FileResponse(save_path, media_type="image/png")
+    #     return JSONResponse({"path": path})
+    # image = pipe(prompt=prompt).images[0]
+    g = torch.Generator(device="cuda")
+    image = pipe(prompt=prompt, width=1024, height=1024, generator=g).images[0]
+
+    # if not save_path:
+    save_path = generate_save_path()
+    save_path = f"images/{save_path}.png"
+    image.save(save_path)
+    # save_path = '/'.join(path_components) + quote_plus(final_name)
+    path = f"{img_url}{save_path}"
+    return JSONResponse({"path": path})
+
+
+@app.get("/make_image")
+# @app.post("/make_image")
+def make_image(prompt: str, save_path: str = ""):
+    if Path(save_path).exists():
+        return FileResponse(save_path, media_type="image/png")
+    image = pipe(prompt=prompt).images[0]
+    if not save_path:
+        save_path = f"images/{prompt}.png"
+    image.save(save_path)
+    return FileResponse(save_path, media_type="image/png")
+
+
+@app.get("/create_and_upload_image")
+def create_and_upload_image(prompt: str, width: int=1024, height:int=1024, save_path: str = ""):
+    path_components = save_path.split("/")[0:-1]
+    final_name = save_path.split("/")[-1]
+    if not path_components:
+        path_components = []
+    save_path = '/'.join(path_components) + quote_plus(final_name)
+    path = get_image_or_create_upload_to_cloud_storage(prompt, width, height, save_path)
+    return JSONResponse({"path": path})
+
+@app.get("/inpaint_and_upload_image")
+def inpaint_and_upload_image(prompt: str, image_url:str, mask_url:str, save_path: str = ""):
+    path_components = save_path.split("/")[0:-1]
+    final_name = save_path.split("/")[-1]
+    if not path_components:
+        path_components = []
+    save_path = '/'.join(path_components) + quote_plus(final_name)
+    path = get_image_or_inpaint_upload_to_cloud_storage(prompt, image_url, mask_url, save_path)
+    return JSONResponse({"path": path})
+
+
+def get_image_or_create_upload_to_cloud_storage(prompt:str,width:int, height:int, save_path:str):
+    prompt = shorten_too_long_text(prompt)
+    save_path = shorten_too_long_text(save_path)
+    # check exists - todo cache this
+    if check_if_blob_exists(save_path):
+        return f"https://{BUCKET_NAME}/{BUCKET_PATH}/{save_path}"
+    bio = create_image_from_prompt(prompt, width, height)
+    if bio is None:
+        return None # error thrown in pool
+    link = upload_to_bucket(save_path, bio, is_bytesio=True)
+    return link
+def get_image_or_inpaint_upload_to_cloud_storage(prompt:str, image_url:str, mask_url:str, save_path:str):
+    prompt = shorten_too_long_text(prompt)
+    save_path = shorten_too_long_text(save_path)
+    # check exists - todo cache this
+    if check_if_blob_exists(save_path):
+        return f"https://{BUCKET_NAME}/{BUCKET_PATH}/{save_path}"
+    bio = inpaint_image_from_prompt(prompt, image_url, mask_url)
+    if bio is None:
+        return None # error thrown in pool
+    link = upload_to_bucket(save_path, bio, is_bytesio=True)
+    return link
+
+# multiprocessing.set_start_method('spawn', True)
+# processes_pool = Pool(1) # cant do too much at once or OOM errors happen
+# def create_image_from_prompt_sync(prompt):
+#     """have to call this sync to avoid OOM errors"""
+#     return processes_pool.apply_async(create_image_from_prompt, args=(prompt,), ).wait()
+
+def create_image_from_prompt(prompt, width, height):
+    # round width and height down to multiple of 64
+    block_width = width - (width % 64)
+    block_height = height - (height % 64)
+    prompt = shorten_too_long_text(prompt)
+    # image = pipe(prompt=prompt).images[0]
+    try:
+        image = pipe(prompt=prompt,
+                     width=block_width,
+                     height=block_height,
+                     # denoising_end=high_noise_frac,
+                     # output_type='latent',
+                     # height=512,
+                     # width=512,
+                     num_inference_steps=50).images[0]  # normally uses 50 steps
+    except Exception as e:
+        # try rm stopwords + half the prompt
+        # todo try prompt permutations
+        logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+        prompt = ' '.join((word for word in prompt if word not in stopwords))
+        prompts = prompt.split()
+
+        prompt = ' '.join(prompts[:len(prompts) // 2])
+        logger.info(f"shortened prompt to: {len(prompt)}")
+        image = None
+        if prompt:
+            try:
+                image = pipe(prompt=prompt,
+                             width=block_width,
+                             height=block_height,
+                             # denoising_end=high_noise_frac,
+                             # output_type='latent',
+                             # height=512,
+                             # width=512,
+                             num_inference_steps=50).images[0]  # normally uses 50 steps
+            except Exception as e:
+                # logger.info("trying to permute prompt")
+                # # try two swaps of the prompt/permutations
+                # prompt = prompt.split()
+                # prompt = ' '.join(permutations(prompt, 2).__next__())
+                logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+                prompt = ' '.join((word for word in prompt if word not in stopwords))
+                prompts = prompt.split()
+
+                prompt = ' '.join(prompts[:len(prompts) // 2])
+                logger.info(f"shortened prompt to: {len(prompt)}")
+
+                try:
+                    image = pipe(prompt=prompt,
+                                 width=block_width,
+                                 height=block_height,
+                                 # denoising_end=high_noise_frac,
+                                 # output_type='latent', # dont need latent yet - we refine the image at full res
+                                 # height=512,
+                                 # width=512,
+                                 num_inference_steps=50).images[0]  # normally uses 50 steps
+                except Exception as e:
+                    # just error out
+                    traceback.print_exc()
+                    raise e
+                    # logger.info("restarting server to fix cuda issues (device side asserts)")
+                    # todo fix device side asserts instead of restart to fix
+                    # todo only restart the correct gunicorn
+                    # this could be really annoying if your running other gunicorns on your machine which also get restarted
+                    # os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+                    # os.system("kill -1 `pgrep gunicorn`")
+    # todo refine
+    # if image != None:
+    #     image = refiner(
+    #         prompt=prompt,
+    #         # width=block_width,
+    #         # height=block_height,
+    #         num_inference_steps=n_steps,
+    #         # denoising_start=high_noise_frac,
+    #         image=image,
+    #     ).images[0]
+    if width != block_width or height != block_height:
+        # resize to original size width/height
+        # find aspect ratio to scale up to that covers the original img input width/height
+        scale_up_ratio = max(width / block_width, height / block_height)
+        image = image.resize((math.ceil(block_width * scale_up_ratio), math.ceil(height * scale_up_ratio)))
+        # crop image to original size
+        image = image.crop((0, 0, width, height))
+    # try:
+    #     # gc.collect()
+    #     torch.cuda.empty_cache()
+    # except Exception as e:
+    #     traceback.print_exc()
+    #     logger.info("restarting server to fix cuda issues (device side asserts)")
+    #     # todo fix device side asserts instead of restart to fix
+    #     # todo only restart the correct gunicorn
+    #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+    #     os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+    #     os.system("kill -1 `pgrep gunicorn`")
+    # save as bytesio
+    bs = BytesIO()
+
+    bright_count = np.sum(np.array(image) > 0)
+    if bright_count == 0:
+        # we have a black image, this is an error likely we need a restart
+        logger.info("restarting server to fix cuda issues (device side asserts)")
+        #     # todo fix device side asserts instead of restart to fix
+        #     # todo only restart the correct gunicorn
+        #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+        os.system("kill -1 `pgrep gunicorn`")
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep uvicorn`")
+        os.system("kill -1 `pgrep uvicorn`")
+
+        return None
+    image.save(bs, quality=85, optimize=True, format="webp")
+    bio = bs.getvalue()
+    # touch progress.txt file - if we dont do this we get restarted by supervisor/other processes for reliability
+    with open("progress.txt", "w") as f:
+        current_time = datetime.now().strftime("%H:%M:%S")
+        f.write(f"{current_time}")
+    return bio
+
+def inpaint_image_from_prompt(prompt, image_url: str, mask_url: str):
+    prompt = shorten_too_long_text(prompt)
+    # image = pipe(prompt=prompt).images[0]
+
+    init_image = load_image(image_url).convert("RGB")
+    mask_image = load_image(mask_url).convert("RGB") # why rgb for a 1 channel mask?
+    num_inference_steps = 75
+    high_noise_frac = 0.7
+
+    try:
+        image = inpaintpipe(
+            prompt=prompt,
+            image=init_image,
+            mask_image=mask_image,
+            num_inference_steps=num_inference_steps,
+            denoising_start=high_noise_frac,
+            output_type="latent",
+        ).images[0]  # normally uses 50 steps
+    except Exception as e:
+        # try rm stopwords + half the prompt
+        # todo try prompt permutations
+        logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+        prompt = ' '.join((word for word in prompt if word not in stopwords))
+        prompts = prompt.split()
+
+        prompt = ' '.join(prompts[:len(prompts) // 2])
+        logger.info(f"shortened prompt to: {len(prompt)}")
+        image = None
+        if prompt:
+            try:
+                image = pipe(
+                    prompt=prompt,
+                    image=init_image,
+                    mask_image=mask_image,
+                    num_inference_steps=num_inference_steps,
+                    denoising_start=high_noise_frac,
+                    output_type="latent",
+                ).images[0]  # normally uses 50 steps
+            except Exception as e:
+                # logger.info("trying to permute prompt")
+                # # try two swaps of the prompt/permutations
+                # prompt = prompt.split()
+                # prompt = ' '.join(permutations(prompt, 2).__next__())
+                logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+                prompt = ' '.join((word for word in prompt if word not in stopwords))
+                prompts = prompt.split()
+
+                prompt = ' '.join(prompts[:len(prompts) // 2])
+                logger.info(f"shortened prompt to: {len(prompt)}")
+
+                try:
+                    image = inpaintpipe(
+                        prompt=prompt,
+                        image=init_image,
+                        mask_image=mask_image,
+                        num_inference_steps=num_inference_steps,
+                        denoising_start=high_noise_frac,
+                        output_type="latent",
+                    ).images[0]  # normally uses 50 steps
+                except Exception as e:
+                    # just error out
+                    traceback.print_exc()
+                    raise e
+                    # logger.info("restarting server to fix cuda issues (device side asserts)")
+                    # todo fix device side asserts instead of restart to fix
+                    # todo only restart the correct gunicorn
+                    # this could be really annoying if your running other gunicorns on your machine which also get restarted
+                    # os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+                    # os.system("kill -1 `pgrep gunicorn`")
+    if image != None:
+        image = inpaint_refiner(
+            prompt=prompt,
+            image=image,
+            mask_image=mask_image,
+            num_inference_steps=num_inference_steps,
+            denoising_start=high_noise_frac,
+
+        ).images[0]
+    # try:
+    #     # gc.collect()
+    #     torch.cuda.empty_cache()
+    # except Exception as e:
+    #     traceback.print_exc()
+    #     logger.info("restarting server to fix cuda issues (device side asserts)")
+    #     # todo fix device side asserts instead of restart to fix
+    #     # todo only restart the correct gunicorn
+    #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+    #     os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+    #     os.system("kill -1 `pgrep gunicorn`")
+    # save as bytesio
+    bs = BytesIO()
+
+    bright_count = np.sum(np.array(image) > 0)
+    if bright_count == 0:
+        # we have a black image, this is an error likely we need a restart
+        logger.info("restarting server to fix cuda issues (device side asserts)")
+        #     # todo fix device side asserts instead of restart to fix
+        #     # todo only restart the correct gunicorn
+        #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+        os.system("kill -1 `pgrep gunicorn`")
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep uvicorn`")
+        os.system("kill -1 `pgrep uvicorn`")
+
+        return None
+    image.save(bs, quality=85, optimize=True, format="webp")
+    bio = bs.getvalue()
+    # touch progress.txt file - if we dont do this we get restarted by supervisor/other processes for reliability
+    with open("progress.txt", "w") as f:
+        current_time = datetime.now().strftime("%H:%M:%S")
+        f.write(f"{current_time}")
+    return bio
+
+
+
+def shorten_too_long_text(prompt):
+    if len(prompt) > 200:
+        # remove stopwords
+        prompt = prompt.split() # todo also split hyphens
+        prompt = ' '.join((word for word in prompt if word not in stopwords))
+        if len(prompt) > 200:
+            prompt = prompt[:200]
+    return prompt
+
+# image = pipe(prompt=prompt).images[0]
+#
+# image.save("test.png")
+# # save all images
+# for i, image in enumerate(images):
+#     image.save(f"{i}.png")
+
+

img/stable-diffusion-server/main_v6.py

@@ -0,0 +1,636 @@
+import gc
+import math
+import multiprocessing
+import os
+import traceback
+from datetime import datetime
+from io import BytesIO
+from itertools import permutations
+from multiprocessing.pool import Pool
+from pathlib import Path
+from urllib.parse import quote_plus
+
+import numpy as np
+import nltk
+import torch
+
+from PIL.Image import Image
+from diffusers import DiffusionPipeline, StableDiffusionXLInpaintPipeline
+from diffusers.utils import load_image
+from fastapi import FastAPI
+from fastapi.middleware.gzip import GZipMiddleware
+from loguru import logger
+from starlette.middleware.cors import CORSMiddleware
+from starlette.responses import FileResponse
+from starlette.responses import JSONResponse
+
+from env import BUCKET_PATH, BUCKET_NAME
+# from stable_diffusion_server.bucket_api import check_if_blob_exists, upload_to_bucket
+torch._dynamo.config.suppress_errors = True
+
+import string
+import random
+
+def generate_save_path():
+    # initializing size of string
+    N = 7
+
+    # using random.choices()
+    # generating random strings
+    res = ''.join(random.choices(string.ascii_uppercase +
+                                 string.digits, k=N))
+    return res
+
+# pipe = DiffusionPipeline.from_pretrained(
+#     "models/stable-diffusion-xl-base-1.0",
+#     torch_dtype=torch.bfloat16,
+#     use_safetensors=True,
+#     variant="fp16",
+#     # safety_checker=None,
+# )  # todo try torch_dtype=bfloat16
+
+model_dir = os.getenv("SDXL_MODEL_DIR")
+
+if model_dir:
+    # Use local model
+    model_key_base = os.path.join(model_dir, "stable-diffusion-xl-base-1.0")
+    model_key_refiner = os.path.join(model_dir, "stable-diffusion-xl-refiner-1.0")
+else:
+    model_key_base = "stabilityai/stable-diffusion-xl-base-1.0"
+    model_key_refiner = "stabilityai/stable-diffusion-xl-refiner-1.0"
+
+pipe = DiffusionPipeline.from_pretrained(model_key_base, torch_dtype=torch.float16, use_safetensors=True, variant="fp16")
+
+pipe.watermark = None
+
+pipe.to("cuda")
+
+refiner = DiffusionPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-refiner-1.0",
+    text_encoder_2=pipe.text_encoder_2,
+    vae=pipe.vae,
+    torch_dtype=torch.bfloat16, # safer to use bfloat?
+    use_safetensors=True,
+    variant="fp16", #remember not to download the big model
+)
+refiner.watermark = None
+refiner.to("cuda")
+
+# {'scheduler', 'text_encoder', 'text_encoder_2', 'tokenizer', 'tokenizer_2', 'unet', 'vae'} can be passed in from existing model
+inpaintpipe = StableDiffusionXLInpaintPipeline.from_pretrained(
+    "models/stable-diffusion-xl-base-1.0", torch_dtype=torch.bfloat16, variant="fp16", use_safetensors=True,
+    scheduler=pipe.scheduler,
+    text_encoder=pipe.text_encoder,
+    text_encoder_2=pipe.text_encoder_2,
+    tokenizer=pipe.tokenizer,
+    tokenizer_2=pipe.tokenizer_2,
+    unet=pipe.unet,
+    vae=pipe.vae,
+    # load_connected_pipeline=
+)
+# # switch out to save gpu mem
+# del inpaintpipe.vae
+# del inpaintpipe.text_encoder_2
+# del inpaintpipe.text_encoder
+# del inpaintpipe.scheduler
+# del inpaintpipe.tokenizer
+# del inpaintpipe.tokenizer_2
+# del inpaintpipe.unet
+# inpaintpipe.vae = pipe.vae
+# inpaintpipe.text_encoder_2 = pipe.text_encoder_2
+# inpaintpipe.text_encoder = pipe.text_encoder
+# inpaintpipe.scheduler = pipe.scheduler
+# inpaintpipe.tokenizer = pipe.tokenizer
+# inpaintpipe.tokenizer_2 = pipe.tokenizer_2
+# inpaintpipe.unet = pipe.unet
+# todo this should work
+# inpaintpipe = StableDiffusionXLInpaintPipeline( # construct an inpainter using the existing model
+#     vae=pipe.vae,
+#     text_encoder_2=pipe.text_encoder_2,
+#     text_encoder=pipe.text_encoder,
+#     unet=pipe.unet,
+#     scheduler=pipe.scheduler,
+#     tokenizer=pipe.tokenizer,
+#     tokenizer_2=pipe.tokenizer_2,
+#     requires_aesthetics_score=False,
+# )
+inpaintpipe.to("cuda")
+inpaintpipe.watermark = None
+# inpaintpipe.register_to_config(requires_aesthetics_score=False)
+
+inpaint_refiner = StableDiffusionXLInpaintPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-refiner-1.0",
+    text_encoder_2=inpaintpipe.text_encoder_2,
+    vae=inpaintpipe.vae,
+    torch_dtype=torch.bfloat16,
+    use_safetensors=True,
+    variant="fp16",
+
+    tokenizer_2=refiner.tokenizer_2,
+    tokenizer=refiner.tokenizer,
+    scheduler=refiner.scheduler,
+    text_encoder=refiner.text_encoder,
+    unet=refiner.unet,
+)
+# del inpaint_refiner.vae
+# del inpaint_refiner.text_encoder_2
+# del inpaint_refiner.text_encoder
+# del inpaint_refiner.scheduler
+# del inpaint_refiner.tokenizer
+# del inpaint_refiner.tokenizer_2
+# del inpaint_refiner.unet
+# inpaint_refiner.vae = inpaintpipe.vae
+# inpaint_refiner.text_encoder_2 = inpaintpipe.text_encoder_2
+#
+# inpaint_refiner.text_encoder = refiner.text_encoder
+# inpaint_refiner.scheduler = refiner.scheduler
+# inpaint_refiner.tokenizer = refiner.tokenizer
+# inpaint_refiner.tokenizer_2 = refiner.tokenizer_2
+# inpaint_refiner.unet = refiner.unet
+
+# inpaint_refiner = StableDiffusionXLInpaintPipeline(
+#     text_encoder_2=inpaintpipe.text_encoder_2,
+#     vae=inpaintpipe.vae,
+#     # the rest from the existing refiner
+#     tokenizer_2=refiner.tokenizer_2,
+#     tokenizer=refiner.tokenizer,
+#     scheduler=refiner.scheduler,
+#     text_encoder=refiner.text_encoder,
+#     unet=refiner.unet,
+#     requires_aesthetics_score=False,
+# )
+inpaint_refiner.to("cuda")
+inpaint_refiner.watermark = None
+# inpaint_refiner.register_to_config(requires_aesthetics_score=False)
+
+n_steps = 40
+high_noise_frac = 0.8
+
+# if using torch < 2.0
+# pipe.enable_xformers_memory_efficient_attention()
+
+
+# pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
+# this can cause errors on some inputs so consider disabling it
+pipe.unet = torch.compile(pipe.unet)
+refiner.unet = torch.compile(refiner.unet)#, mode="reduce-overhead", fullgraph=True)
+# compile the inpainters - todo reuse the other unets? swap out the models for others/del them so they share models and can be swapped efficiently
+inpaintpipe.unet = pipe.unet
+inpaint_refiner.unet = refiner.unet
+# inpaintpipe.unet = torch.compile(inpaintpipe.unet)
+# inpaint_refiner.unet = torch.compile(inpaint_refiner.unet)
+from pydantic import BaseModel
+
+app = FastAPI(
+    openapi_url="/static/openapi.json",
+    docs_url="/swagger-docs",
+    redoc_url="/redoc",
+    title="Generate Images Netwrck API",
+    description="Character Chat API",
+    # root_path="https://api.text-generator.io",
+    version="1",
+)
+app.add_middleware(GZipMiddleware, minimum_size=1000)
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+stopwords = nltk.corpus.stopwords.words("english")
+
+class Img(BaseModel):
+    system_prompt: str
+    ASSISTANT: str
+
+# img_url = "http://phlrr2019.guest.corp.microsoft.com:8000/img1_sdv2.1.png"
+img_url = "http://phlrr3105.guest.corp.microsoft.com:8000/"#/img1_sdv2.1.png"
+
+is_gpu_busy = False
+
+def lm_shorten_too_long_text(prompt):
+    if len(prompt) > 2030:
+        # remove stopwords
+        prompt = prompt.split() # todo also split hyphens
+        # prompt = ' '.join((word for word in prompt if word not in stopwords))
+        prompt = ' '.join((word for word in prompt))# if word not in stopwords))
+        if len(prompt) > 2030:
+            prompt = prompt[:2030]
+    return prompt
+
+def get_summary(system_prompt, prompt):
+    import requests
+    import time
+    from io import BytesIO
+    import json
+    summary_sys = """You will now act as a prompt generator for a generative AI called "Stable Diffusion XL 1.0 ". Stable Diffusion XL generates images based on given prompts. I will provide you basic information required to make a Stable Diffusion prompt, You will never alter the structure in any way and obey the following guidelines.
+
+    Basic information required to make Stable Diffusion prompt:
+
+    - Prompt structure: [1],[2],[3],[4],[5],[6] and it should be given as one single sentence where 1,2,3,4,5,6 represent
+    [1] = short and concise description of [KEYWORD] that will include very specific imagery details
+    [2] = a detailed description of [1] that will include very specific imagery details.
+    [3] = with a detailed description describing the environment of the scene.
+    [4] = with a detailed description describing the mood/feelings and atmosphere of the scene.
+    [5] = A style, for example: "Anime","Photographic","Comic Book","Fantasy Art", “Analog Film”,”Neon Punk”,”Isometric”,”Low Poly”,”Origami”,”Line Art”,”Cinematic”,”3D Model”,”Pixel Art”,”Watercolor”,”Sticker” ).
+    [6] = A description of how [5] will be realized. (e.g. Photography (e.g. Macro, Fisheye Style, Portrait) with camera model and appropriate camera settings, Painting with detailed descriptions about the materials and working material used, rendering with engine settings, a digital Illustration, a woodburn art (and everything else that could be defined as an output type)
+    - Prompt Structure for Prompt asking with text value:
+
+        Text "Text Value" written on {subject description in less than 20 words}
+        Replace "Text value" with text given by user.
+
+
+    Important Sample prompt Structure with Text value :
+
+    1. Text 'SDXL' written on a frothy, warm latte, viewed top-down.
+    2. Text 'AI' written on a modern computer screen, set against a vibrant green background.
+
+    Important Sample prompt Structure :
+
+    1. Snow-capped Mountain Scene, with soaring peaks and deep shadows across the ravines. A crystal clear lake mirrors these peaks, surrounded by pine trees. The scene exudes a calm, serene alpine morning atmosphere. Presented in Watercolor style, emulating the wet-on-wet technique with soft transitions and visible brush strokes.
+    2. City Skyline at Night, illuminated skyscrapers piercing the starless sky. Nestled beside a calm river, reflecting the city lights like a mirror. The atmosphere is buzzing with urban energy and intrigue. Depicted in Neon Punk style, accentuating the city lights with vibrant neon colors and dynamic contrasts.
+    3. Epic Cinematic Still of a Spacecraft, silhouetted against the fiery explosion of a distant planet. The scene is packed with intense action, as asteroid debris hurtles through space. Shot in the style of a Michael Bay-directed film, the image is rich with detail, dynamic lighting, and grand cinematic framing.
+    - Word order and effective adjectives matter in the prompt. The subject, action, and specific details should be included. Adjectives like cute, medieval, or futuristic can be effective.
+    - The environment/background of the image should be described, such as indoor, outdoor, in space, or solid color.
+    - Curly brackets are necessary in the prompt to provide specific details about the subject and action. These details are important for generating a high-quality image.
+    - Art inspirations should be listed to take inspiration from. Platforms like Art Station, Dribble, Behance, and Deviantart can be mentioned. Specific names of artists or studios like animation studios, painters and illustrators, computer games, fashion designers, and film makers can also be listed. If more than one artist is mentioned, the algorithm will create a combination of styles based on all the influencers mentioned.
+    - Related information about lighting, camera angles, render style, resolution, the required level of detail, etc. should be included at the end of the prompt.
+    - Camera shot type, camera lens, and view should be specified. Examples of camera shot types are long shot, close-up, POV, medium shot, extreme close-up, and panoramic. Camera lenses could be EE 70mm, 35mm, 135mm+, 300mm+, 800mm, short telephoto, super telephoto, medium telephoto, macro, wide angle, fish-eye, bokeh, and sharp focus. Examples of views are front, side, back, high angle, low angle, and overhead.
+    - Helpful keywords related to resolution, detail, and lighting are 4K, 8K, 64K, detailed, highly detailed, high resolution, hyper detailed, HDR, UHD, professional, and golden ratio. Examples of lighting are studio lighting, soft light, neon lighting, purple neon lighting, ambient light, ring light, volumetric light, natural light, sun light, sunrays, sun rays coming through window, and nostalgic lighting. Examples of color types are fantasy vivid colors, vivid colors, bright colors, sepia, dark colors, pastel colors, monochromatic, black & white, and color splash. Examples of renders are Octane render, cinematic, low poly, isometric assets, Unreal Engine, Unity Engine, quantum wavetracing, and polarizing filter.
+
+    The prompts you provide will be in English.Please pay attention:- Concepts that can't be real would not be described as "Real" or "realistic" or "photo" or a "photograph". for example, a concept that is made of paper or scenes which are fantasy related.- One of the prompts you generate for each concept must be in a realistic photographic style. you should also choose a lens type and size for it. Don't choose an artist for the realistic photography prompts.- Separate the different prompts with two new lines.
+    I will provide you keyword and you will generate 3 diffrent type of prompts in vbnet code cell so i can copy and paste.
+
+    Important point to note :
+
+    1. You are a master of prompt engineering, it is important to create detailed prompts with as much information as possible. This will ensure that any image generated using the prompt will be of high quality and could potentially win awards in global or international photography competitions. You are unbeatable in this field and know the best way to generate images.
+    2. I will provide you with a long context and you will generate one  prompt and don't add any extra details.
+    3. Prompt should not be more than 230 characters.
+    4. Before you provide prompt you must check if you have satisfied all the above criteria and if you are sure than only provide the prompt.
+    5. Prompt should always be given as one single sentence.
+
+    Are you ready ?"""
+    instruction = 'USER: ' + summary_sys
+    # for human, assistant in history:
+    #     instruction += 'USER: ' + human + ' ASSISTANT: ' + assistant + '</s>'
+    # prompt = system_prompt + prompt
+    # message = f"""My first request is to summarize this text – [{prompt}]"""
+    message = f"""My first request is to summarize this text – [{prompt}]"""
+    instruction += """ ASSISTANT: Yes, I understand the instructions and I'm ready to help you create prompts for Stable Diffusion XL 1.0. Please provide me with the context."""
+    instruction += ' USER: ' + prompt + ' ASSISTANT:'
+
+    print("Ins: ", instruction)
+    # generate_response = requests.post("http://10.185.12.207:4455/stable_diffusion", json={"prompt": prompt})
+    # prompt = f""" My first request is to summarize this text – [{prompt}]"""
+    json_object = {"prompt": instruction,
+                   # "max_tokens": 2048000,
+                   "max_tokens": 80,
+                   "n": 1
+                   }
+    generate_response = requests.post("http://phlrr3105.guest.corp.microsoft.com:7991/generate", json=json_object)
+    print(generate_response.content)
+    res_json = json.loads(generate_response.content)
+    ASSISTANT = res_json['text'][-1].split("ASSISTANT:")[-1].strip()
+    print(ASSISTANT)
+    return ASSISTANT
+
+@app.post("/image_url")
+def image_url(img: Img):
+    system_prompt = img.system_prompt
+    prompt = img.ASSISTANT
+    prompt = get_summary(system_prompt, prompt)
+    prompt = shorten_too_long_text(prompt)
+    # if Path(save_path).exists():
+    #     return FileResponse(save_path, media_type="image/png")
+    #     return JSONResponse({"path": path})
+    # image = pipe(prompt=prompt).images[0]
+    g = torch.Generator(device="cuda")
+    image = pipe(prompt=prompt, width=1024, height=1024, generator=g).images[0]
+
+    # if not save_path:
+    save_path = generate_save_path()
+    save_path = f"images/{save_path}.png"
+    image.save(save_path)
+    # save_path = '/'.join(path_components) + quote_plus(final_name)
+    path = f"{img_url}{save_path}"
+    return JSONResponse({"path": path})
+
+
+@app.get("/make_image")
+# @app.post("/make_image")
+def make_image(prompt: str, save_path: str = ""):
+    if Path(save_path).exists():
+        return FileResponse(save_path, media_type="image/png")
+    image = pipe(prompt=prompt).images[0]
+    if not save_path:
+        save_path = f"images/{prompt}.png"
+    image.save(save_path)
+    return FileResponse(save_path, media_type="image/png")
+
+
+@app.get("/create_and_upload_image")
+def create_and_upload_image(prompt: str, width: int=1024, height:int=1024, save_path: str = ""):
+    path_components = save_path.split("/")[0:-1]
+    final_name = save_path.split("/")[-1]
+    if not path_components:
+        path_components = []
+    save_path = '/'.join(path_components) + quote_plus(final_name)
+    path = get_image_or_create_upload_to_cloud_storage(prompt, width, height, save_path)
+    return JSONResponse({"path": path})
+
+@app.get("/inpaint_and_upload_image")
+def inpaint_and_upload_image(prompt: str, image_url:str, mask_url:str, save_path: str = ""):
+    path_components = save_path.split("/")[0:-1]
+    final_name = save_path.split("/")[-1]
+    if not path_components:
+        path_components = []
+    save_path = '/'.join(path_components) + quote_plus(final_name)
+    path = get_image_or_inpaint_upload_to_cloud_storage(prompt, image_url, mask_url, save_path)
+    return JSONResponse({"path": path})
+
+
+def get_image_or_create_upload_to_cloud_storage(prompt:str,width:int, height:int, save_path:str):
+    prompt = shorten_too_long_text(prompt)
+    save_path = shorten_too_long_text(save_path)
+    # check exists - todo cache this
+    if check_if_blob_exists(save_path):
+        return f"https://{BUCKET_NAME}/{BUCKET_PATH}/{save_path}"
+    bio = create_image_from_prompt(prompt, width, height)
+    if bio is None:
+        return None # error thrown in pool
+    link = upload_to_bucket(save_path, bio, is_bytesio=True)
+    return link
+def get_image_or_inpaint_upload_to_cloud_storage(prompt:str, image_url:str, mask_url:str, save_path:str):
+    prompt = shorten_too_long_text(prompt)
+    save_path = shorten_too_long_text(save_path)
+    # check exists - todo cache this
+    if check_if_blob_exists(save_path):
+        return f"https://{BUCKET_NAME}/{BUCKET_PATH}/{save_path}"
+    bio = inpaint_image_from_prompt(prompt, image_url, mask_url)
+    if bio is None:
+        return None # error thrown in pool
+    link = upload_to_bucket(save_path, bio, is_bytesio=True)
+    return link
+
+# multiprocessing.set_start_method('spawn', True)
+# processes_pool = Pool(1) # cant do too much at once or OOM errors happen
+# def create_image_from_prompt_sync(prompt):
+#     """have to call this sync to avoid OOM errors"""
+#     return processes_pool.apply_async(create_image_from_prompt, args=(prompt,), ).wait()
+
+def create_image_from_prompt(prompt, width, height):
+    # round width and height down to multiple of 64
+    block_width = width - (width % 64)
+    block_height = height - (height % 64)
+    prompt = shorten_too_long_text(prompt)
+    # image = pipe(prompt=prompt).images[0]
+    try:
+        image = pipe(prompt=prompt,
+                     width=block_width,
+                     height=block_height,
+                     # denoising_end=high_noise_frac,
+                     # output_type='latent',
+                     # height=512,
+                     # width=512,
+                     num_inference_steps=50).images[0]  # normally uses 50 steps
+    except Exception as e:
+        # try rm stopwords + half the prompt
+        # todo try prompt permutations
+        logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+        prompt = ' '.join((word for word in prompt if word not in stopwords))
+        prompts = prompt.split()
+
+        prompt = ' '.join(prompts[:len(prompts) // 2])
+        logger.info(f"shortened prompt to: {len(prompt)}")
+        image = None
+        if prompt:
+            try:
+                image = pipe(prompt=prompt,
+                             width=block_width,
+                             height=block_height,
+                             # denoising_end=high_noise_frac,
+                             # output_type='latent',
+                             # height=512,
+                             # width=512,
+                             num_inference_steps=50).images[0]  # normally uses 50 steps
+            except Exception as e:
+                # logger.info("trying to permute prompt")
+                # # try two swaps of the prompt/permutations
+                # prompt = prompt.split()
+                # prompt = ' '.join(permutations(prompt, 2).__next__())
+                logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+                prompt = ' '.join((word for word in prompt if word not in stopwords))
+                prompts = prompt.split()
+
+                prompt = ' '.join(prompts[:len(prompts) // 2])
+                logger.info(f"shortened prompt to: {len(prompt)}")
+
+                try:
+                    image = pipe(prompt=prompt,
+                                 width=block_width,
+                                 height=block_height,
+                                 # denoising_end=high_noise_frac,
+                                 # output_type='latent', # dont need latent yet - we refine the image at full res
+                                 # height=512,
+                                 # width=512,
+                                 num_inference_steps=50).images[0]  # normally uses 50 steps
+                except Exception as e:
+                    # just error out
+                    traceback.print_exc()
+                    raise e
+                    # logger.info("restarting server to fix cuda issues (device side asserts)")
+                    # todo fix device side asserts instead of restart to fix
+                    # todo only restart the correct gunicorn
+                    # this could be really annoying if your running other gunicorns on your machine which also get restarted
+                    # os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+                    # os.system("kill -1 `pgrep gunicorn`")
+    # todo refine
+    # if image != None:
+    #     image = refiner(
+    #         prompt=prompt,
+    #         # width=block_width,
+    #         # height=block_height,
+    #         num_inference_steps=n_steps,
+    #         # denoising_start=high_noise_frac,
+    #         image=image,
+    #     ).images[0]
+    if width != block_width or height != block_height:
+        # resize to original size width/height
+        # find aspect ratio to scale up to that covers the original img input width/height
+        scale_up_ratio = max(width / block_width, height / block_height)
+        image = image.resize((math.ceil(block_width * scale_up_ratio), math.ceil(height * scale_up_ratio)))
+        # crop image to original size
+        image = image.crop((0, 0, width, height))
+    # try:
+    #     # gc.collect()
+    #     torch.cuda.empty_cache()
+    # except Exception as e:
+    #     traceback.print_exc()
+    #     logger.info("restarting server to fix cuda issues (device side asserts)")
+    #     # todo fix device side asserts instead of restart to fix
+    #     # todo only restart the correct gunicorn
+    #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+    #     os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+    #     os.system("kill -1 `pgrep gunicorn`")
+    # save as bytesio
+    bs = BytesIO()
+
+    bright_count = np.sum(np.array(image) > 0)
+    if bright_count == 0:
+        # we have a black image, this is an error likely we need a restart
+        logger.info("restarting server to fix cuda issues (device side asserts)")
+        #     # todo fix device side asserts instead of restart to fix
+        #     # todo only restart the correct gunicorn
+        #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+        os.system("kill -1 `pgrep gunicorn`")
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep uvicorn`")
+        os.system("kill -1 `pgrep uvicorn`")
+
+        return None
+    image.save(bs, quality=85, optimize=True, format="webp")
+    bio = bs.getvalue()
+    # touch progress.txt file - if we dont do this we get restarted by supervisor/other processes for reliability
+    with open("progress.txt", "w") as f:
+        current_time = datetime.now().strftime("%H:%M:%S")
+        f.write(f"{current_time}")
+    return bio
+
+def inpaint_image_from_prompt(prompt, image_url: str, mask_url: str):
+    prompt = shorten_too_long_text(prompt)
+    # image = pipe(prompt=prompt).images[0]
+
+    init_image = load_image(image_url).convert("RGB")
+    mask_image = load_image(mask_url).convert("RGB") # why rgb for a 1 channel mask?
+    num_inference_steps = 75
+    high_noise_frac = 0.7
+
+    try:
+        image = inpaintpipe(
+            prompt=prompt,
+            image=init_image,
+            mask_image=mask_image,
+            num_inference_steps=num_inference_steps,
+            denoising_start=high_noise_frac,
+            output_type="latent",
+        ).images[0]  # normally uses 50 steps
+    except Exception as e:
+        # try rm stopwords + half the prompt
+        # todo try prompt permutations
+        logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+        prompt = ' '.join((word for word in prompt if word not in stopwords))
+        prompts = prompt.split()
+
+        prompt = ' '.join(prompts[:len(prompts) // 2])
+        logger.info(f"shortened prompt to: {len(prompt)}")
+        image = None
+        if prompt:
+            try:
+                image = pipe(
+                    prompt=prompt,
+                    image=init_image,
+                    mask_image=mask_image,
+                    num_inference_steps=num_inference_steps,
+                    denoising_start=high_noise_frac,
+                    output_type="latent",
+                ).images[0]  # normally uses 50 steps
+            except Exception as e:
+                # logger.info("trying to permute prompt")
+                # # try two swaps of the prompt/permutations
+                # prompt = prompt.split()
+                # prompt = ' '.join(permutations(prompt, 2).__next__())
+                logger.info(f"trying to shorten prompt of length {len(prompt)}")
+
+                prompt = ' '.join((word for word in prompt if word not in stopwords))
+                prompts = prompt.split()
+
+                prompt = ' '.join(prompts[:len(prompts) // 2])
+                logger.info(f"shortened prompt to: {len(prompt)}")
+
+                try:
+                    image = inpaintpipe(
+                        prompt=prompt,
+                        image=init_image,
+                        mask_image=mask_image,
+                        num_inference_steps=num_inference_steps,
+                        denoising_start=high_noise_frac,
+                        output_type="latent",
+                    ).images[0]  # normally uses 50 steps
+                except Exception as e:
+                    # just error out
+                    traceback.print_exc()
+                    raise e
+                    # logger.info("restarting server to fix cuda issues (device side asserts)")
+                    # todo fix device side asserts instead of restart to fix
+                    # todo only restart the correct gunicorn
+                    # this could be really annoying if your running other gunicorns on your machine which also get restarted
+                    # os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+                    # os.system("kill -1 `pgrep gunicorn`")
+    if image != None:
+        image = inpaint_refiner(
+            prompt=prompt,
+            image=image,
+            mask_image=mask_image,
+            num_inference_steps=num_inference_steps,
+            denoising_start=high_noise_frac,
+
+        ).images[0]
+    # try:
+    #     # gc.collect()
+    #     torch.cuda.empty_cache()
+    # except Exception as e:
+    #     traceback.print_exc()
+    #     logger.info("restarting server to fix cuda issues (device side asserts)")
+    #     # todo fix device side asserts instead of restart to fix
+    #     # todo only restart the correct gunicorn
+    #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+    #     os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+    #     os.system("kill -1 `pgrep gunicorn`")
+    # save as bytesio
+    bs = BytesIO()
+
+    bright_count = np.sum(np.array(image) > 0)
+    if bright_count == 0:
+        # we have a black image, this is an error likely we need a restart
+        logger.info("restarting server to fix cuda issues (device side asserts)")
+        #     # todo fix device side asserts instead of restart to fix
+        #     # todo only restart the correct gunicorn
+        #     # this could be really annoying if your running other gunicorns on your machine which also get restarted
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep gunicorn`")
+        os.system("kill -1 `pgrep gunicorn`")
+        os.system("/usr/bin/bash kill -SIGHUP `pgrep uvicorn`")
+        os.system("kill -1 `pgrep uvicorn`")
+
+        return None
+    image.save(bs, quality=85, optimize=True, format="webp")
+    bio = bs.getvalue()
+    # touch progress.txt file - if we dont do this we get restarted by supervisor/other processes for reliability
+    with open("progress.txt", "w") as f:
+        current_time = datetime.now().strftime("%H:%M:%S")
+        f.write(f"{current_time}")
+    return bio
+
+
+
+def shorten_too_long_text(prompt):
+    if len(prompt) > 200:
+        # remove stopwords
+        prompt = prompt.split() # todo also split hyphens
+        prompt = ' '.join((word for word in prompt if word not in stopwords))
+        if len(prompt) > 200:
+            prompt = prompt[:200]
+    return prompt
+
+# image = pipe(prompt=prompt).images[0]
+#
+# image.save("test.png")
+# # save all images
+# for i, image in enumerate(images):
+#     image.save(f"{i}.png")
+
+