INIT

runner.py

@@ -0,0 +1,131 @@
+import sys
+import os
+import uuid
+import json
+import logging
+import traceback
+from datetime import datetime as dt
+
+import torch
+from diffusers import DiffusionPipeline
+
+# Setup Logging
+logging.basicConfig(
+    level=logging.DEBUG,
+    # level=logging.INFO,
+    format="%(asctime)s [%(levelname)s] %(message)s",
+    handlers=[
+        logging.FileHandler("/var/log/" + str(dt.today().strftime('%Y-%m-%d')) + "_-_cron.log"),
+        logging.StreamHandler(sys.stdout)
+    ]
+)
+
+try:
+    for root, dirs, files in os.walk('/app/tasks'):
+        for file in files:
+            try:
+                with open(os.path.join(root, file)) as f:
+                    d = json.load(f)
+                    print(json.dumps(d, indent=4, sort_keys=True))
+
+                    prompt = d["prompt"]
+                    neg_prompt = d["neg_prompt"]
+                    style = d["style"]
+
+                    if "iterations" in d:
+                        iterations = int(d["iterations"])
+                    else:
+                        iterations = 40
+
+                    if "format" in d:
+                        format = d["format"]
+                    else:
+                        format = "square"
+
+                    match format.lower():
+                        case "portrait":
+                            format_width = 915
+                            format_height = 1144
+                        case "widescreen":
+                            format_width = 1344
+                            format_height = 768
+                        case "photo":  # Photo (4x3)
+                            format_width = 1182
+                            format_height = 886
+                        case "square":
+                            format_width = 1024
+                            format_height = 1024
+
+                    # For Debugging decrease image size
+                    # scale_factor = 0.70   # Error: Image shredded
+                    # scale_factor = 0.68   # Error: Image shredded
+                    scale_factor = 0.65  # Working on MacBook Pro (M3 - Nov. 2023)
+
+                    format_width = int(round(round(format_width * scale_factor) / 8) * 8)
+                    format_height = int(round(round(format_height * scale_factor) / 8) * 8)
+
+                with open('/app/styles.json') as f:
+                    for entry in json.load(f):
+                        if entry['style'] == style:
+                            prompt += ", " + entry['prompt_extension']
+                            neg_prompt += "extra fingers, mutated hands, poorly drawn hands, poorly drawn face, deformed, ugly, blurry, bad anatomy, bad proportions, extra limbs, glitchy, double torso, extra arms, extra hands, mangled fingers, missing lips, ugly face, distorted face, extra legs"
+                            neg_prompt += entry['negative_prompt']
+
+                base = DiffusionPipeline.from_pretrained(
+                    "RunDiffusion/Juggernaut-XL-v8"
+                ).to("cpu")
+                refiner = DiffusionPipeline.from_pretrained(
+                    "RunDiffusion/Juggernaut-XL-v8",
+                    text_encoder_2=base.text_encoder_2,
+                    vae=base.vae,
+                ).to("cpu")
+
+                n_steps = iterations
+                high_noise_frac = 0.8
+
+                image = base(
+                    prompt=prompt,
+                    generator=torch.Generator(device="cpu").manual_seed(torch.Generator(device="cpu").seed()),
+                    # generator=torch.Generator(device="cpu").manual_seed(420),
+                    negative_prompt=neg_prompt,
+                    num_inference_steps=n_steps,
+                    denoising_end=high_noise_frac,
+                    width=format_width,
+                    height=format_height,
+                    original_size=(format_width, format_height),
+                    target_size=(format_width, format_height),
+                    output_type="latent",
+                ).images
+                image = refiner(
+                    prompt=prompt,
+                    generator=torch.Generator(device="cpu").manual_seed(torch.Generator(device="cpu").seed()),
+                    # generator=torch.Generator(device="cpu").manual_seed(420),
+                    negative_prompt=neg_prompt,
+                    num_inference_steps=n_steps,
+                    denoising_start=high_noise_frac,
+                    width=format_width,
+                    height=format_height,
+                    original_size=(format_width, format_height),
+                    target_size=(format_width, format_height),
+                    image=image,
+                ).images[0]
+
+                img_uuid = str(uuid.uuid4())
+                file_name = '/app/images/' + img_uuid + '.png'
+                image.save(file_name)
+
+                logging.debug("Image generated and written to file: " + file_name)
+                sys.exit(0)
+
+            except Exception as e:
+                logging.debug("There was an error: " + str(e))
+                logging.debug("Stacktrace: " + str(traceback.format_exc()))
+                logging.debug("Skipping iteration!!!")
+                continue
+
+except Exception as e:
+    logging.debug("There was an error: " + str(e))
+    logging.debug("Stacktrace: " + str(traceback.format_exc()))
+
+finally:
+    exit(0)