Sentdex · September 16, 2022 01:38 · Sep 16, 2022
diff --git a/imagery_create.py b/imagery_create.py
@@ -0,0 +1,57 @@
+import torch
+from diffusers import StableDiffusionPipeline
+from torch import autocast
+import random
+import matplotlib.pyplot as plt
+import os
+
+
+prompts = [
+          "1965 Porsche 911",
+          "1975 Porsche 911",
+          "1985 Porsche 911",
+          "1995 Porsche 911",
+          "2005 Porsche 911 front",
+          "2015 Porsche 911",
+          "2020 Porsche 911",
+          "2020 Porsche 911 GT3 RS"]
+
+
+# make sure you're logged in with `huggingface-cli login`
+pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", 
+                                                revision="fp16",
+                                                torch_dtype=torch.float16, 
+                                                use_auth_token=True)  
+pipe = pipe.to("cuda")
+
+def infer(prompt, num_inference_steps=50, 
+            samples=5, seed=1024, guidance_scale=7.5, 
+            width=512, height=512):
+    generator = torch.Generator("cuda").manual_seed(seed)
+    w = width//8*8
+    h = height//8*8
+    with autocast("cuda"):
+        image = pipe(prompt, guidance_scale=7.5, 
+                    generator=generator, width=w, height=h, 
+                    num_inference_steps=num_inference_steps)["sample"][0]
+    return image
+
+
+for p in prompts:
+
+    prompt_orig = p 
+    if not os.path.exists("imagery"):
+        os.mkdir("imagery")
+
+    if not os.path.exists(f"imagery/{prompt_orig}"):
+        os.mkdir(f"imagery/{prompt_orig}/")
+
+    HM = 200
+    for i in range(HM):
+        print(f"{i+1}/{HM}")
+        prompt_to_use = prompt_orig
+        seed = random.randint(0, 10000)
+        print(seed)
+        image = infer(prompt_to_use, num_inference_steps=75, 
+                      seed=seed, width=512, height=512)
+        image.save(f"imagery/{prompt_to_use}/{seed}.png")