main created

main.py

@@ -0,0 +1,9 @@
+from click import prompt
+from nvidia import execute_nvidia  
+
+
+prompt = "ezreal"
+filename = "ezreal22"
+
+
+execute_nvidia(prompt,filename)

nvidia.py

@@ -1,14 +1,16 @@
 from diffusers import StableDiffusionPipeline
 
-DEVICE = 'cuda'
-
-pipe = StableDiffusionPipeline.from_pretrained(
-    "CompVis/stable-diffusion-v1-4",
-    use_auth_token=True
-).to(DEVICE)
-prompt = "Elon Musk riding elephant on the moon."
-pipe.enable_attention_slicing()
-# _ = pipe(prompt, num_inference_steps=1)
-image = pipe(prompt).images[0]
-image.save("pictures/elon_musk.png")
+def execute_nvidia(prompt, filename):
+
+    DEVICE = 'cuda'
+
+    pipe = StableDiffusionPipeline.from_pretrained(
+        "CompVis/stable-diffusion-v1-4",
+        use_auth_token=True
+    ).to(DEVICE)
+    # prompt = "Elon Musk riding elephant on the moon."
+    pipe.enable_attention_slicing()
+    # _ = pipe(prompt, num_inference_steps=1)
+    image = pipe(prompt).images[0]
+    image.save("pictures/"+ filename +".png")