initial code commit functioning flux script

flux-schnell.py

@@ -0,0 +1,52 @@
+from diffusers import FluxPipeline
+import torch
+from pathlib import Path
+import re
+from datetime import datetime
+
+def slugify(text):
+    # remove non-word characters and foreign characters
+    text = re.sub(r"[^\w\s]", "", text)
+    text = re.sub(r"\s+", "-", text)
+    return text
+
+prompt = "the town center of a small futuristic town with a fountain and geometric dodecahedral d20 buildings with a road leading out to grassy plains dotted with very tall trees"
+height, width = 720, 1280
+
+ckpt_id = "black-forest-labs/FLUX.1-schnell"
+
+DIR_NAME="/nas/dockerdata/immich/local/"
+dirpath = Path(DIR_NAME)
+# create parent dir if doesn't exist
+dirpath.mkdir(parents=True, exist_ok=True)
+
+# denoising
+pipe = FluxPipeline.from_pretrained(
+    ckpt_id,
+    torch_dtype=torch.bfloat16,
+    use_safetensors=True,
+)
+pipe.vae.enable_tiling()
+pipe.vae.enable_slicing()
+pipe.enable_sequential_cpu_offload() # offloads modules to CPU on a submodule level (rather than model level)
+
+output = pipe(
+    prompt,
+    height=height,
+    width=width,
+    num_images_per_prompt=4,
+    num_inference_steps=4,
+    max_sequence_length=256,
+    guidance_scale=0.0,
+)
+print('Max mem allocated (GB) while denoising:', torch.cuda.max_memory_allocated() / (1024 ** 3))
+
+# import matplotlib.pyplot as plt
+# plt.imshow(image)
+# image.save("./whitehenge.png")
+# plt.show()
+for idx, image in enumerate(output.images):
+    timestamp = datetime.now().strftime("%Y%m%d%-H%M%S")
+    image_name = f'{slugify(prompt)}-{idx}-{timestamp}.png'
+    image_path = dirpath / image_name
+    image.save(image_path)

img2img.py

@@ -0,0 +1,40 @@
+from diffusers import StableDiffusionPipeline, EulerDiscreteScheduler
+import torch
+from pathlib import Path
+import re
+
+def slugify(text):
+    # remove non-word characters and foreign characters
+    text = re.sub(r"[^\w\s]", "", text)
+    text = re.sub(r"\s+", "-", text)
+    return text
+
+model_id = "stabilityai/stable-diffusion-2"
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+# Use the Euler scheduler here instead
+scheduler = EulerDiscreteScheduler.from_pretrained(
+    model_id, subfolder="scheduler")
+pipe = StableDiffusionPipeline.from_pretrained(
+    model_id, scheduler=scheduler, torch_dtype=torch.float16)
+pipe = pipe.to(device)
+
+DIR_NAME="./images/"
+dirpath = Path(DIR_NAME)
+# create parent dir if doesn't exist
+dirpath.mkdir(parents=True, exist_ok=True)
+
+prompt = "a photo of an astronaut riding a horse on mars"
+negative_prompt = "blurry, dark photo, blue"
+steps = 15
+scale = 9
+num_images_per_prompt = 5
+seed = torch.randint(0, 1000000, (1,)).item()
+generator = torch.Generator(device=device).manual_seed(seed)
+output = pipe(prompt, negative_prompt=negative_prompt, width=512, height=512, num_inference_steps=steps,
+             guidance_scale=scale, num_images_per_prompt=num_images_per_prompt, generator=generator)
+
+for idx, image in enumerate(output.images):
+    image_name = f'{slugify(prompt)}-{idx}.png'
+    image_path = dirpath / image_name
+    image.save(image_path)

initimg2img.py

@@ -0,0 +1,55 @@
+from diffusers import StableDiffusionImg2ImgPipeline, EulerDiscreteScheduler
+from pathlib import Path
+from PIL import Image
+import torch
+import re
+import requests
+
+
+def slugify(text):
+    # remove non-word characters and foreign characters
+    text = re.sub(r"[^\w\s]", "", text)
+    text = re.sub(r"\s+", "-", text)
+    return text
+
+model_id = "stabilityai/stable-diffusion-2"
+
+images_url = ["https://s3.amazonaws.com/moonup/production/uploads/1675140495576-noauth.png",
+              "https://s3.amazonaws.com/moonup/production/uploads/1675032939263-noauth.png",
+              "https://s3.amazonaws.com/moonup/production/uploads/1673856328001-noauth.png"]
+            
+init_images = [Image.open(requests.get(url, stream=True).raw).convert("RGB").resize((758,768)) for url in images_url]
+
+prompts = ["beautiful colorful flowr",
+           "green city future mountain 3d sunrise skycrapers", 
+           "rainbow beach, palm trees, neon, miami"]
+
+negative_prompts = ["blurry, dark photo, blue",
+                    "blurry, dark photo, blue",
+                    "blurry, dark photo, blue"]
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+# Use the Euler scheduler here instead
+scheduler = EulerDiscreteScheduler.from_pretrained(
+    model_id, subfolder="scheduler")
+pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
+    model_id, scheduler=scheduler, torch_dtype=torch.float16)
+pipe = pipe.to(device)
+
+DIR_NAME="./images/"
+dirpath = Path(DIR_NAME)
+# create parent dir if doesn't exist
+dirpath.mkdir(parents=True, exist_ok=True)
+
+steps = 20
+scale = 9
+num_images_per_prompt = 1
+seed = torch.randint(0, 1000000, (1,)).item()
+generator = torch.Generator(device=device).manual_seed(seed)
+output = pipe(prompts, negative_prompt=negative_prompts, image=init_images, num_inference_steps=steps,
+             guidance_scale=scale, num_images_per_prompt=num_images_per_prompt, generator=generator)
+
+for idx, (image,prompt) in enumerate(zip(output.images, prompts*num_images_per_prompt)):
+    image_name = f'{slugify(prompt)}-{idx}.png'
+    image_path = dirpath / image_name
+    image.save(image_path)