Support for 6gb cards

README.md

@@ -24,8 +24,9 @@ c. Create a write token/Copy an existing Token key and enter in the cmd window.
 5. You can find output files in .\\data\\out\\%Timestamp% folder. Jpeg contain commentaries in metadata with some settings used while generating.
 		
 ## Requirements
-* 11GB of free space.
-* Nvidia card with 8GB+ video memory.
-	* Currently min dalle can run in 6gb. Testing purpose
-## Additional info
-Tested on RTX3070. One picture was making 12 - 14 seconds.
+* 10GB of free space.
+* Nvidia card with 6GB+ video memory.
+
+## Tests
+RTX3070:	 	12 - 14 seconds per picture.
+RTX2060_Laptop:	21 - 25 seconds per picture.

scripts/stable_diff.py

@@ -20,37 +20,38 @@ from diffusers import (
 from diffusers import StableDiffusionImg2ImgPipeline 
 
 t = torch.cuda.get_device_properties(0).total_memory
-if t <=8500000000:
-    print("Not enough GPU memory to generate pictures")
-else:
-    file1 = open("prompt.txt","r+")
-    text = file1.read()
-    print(text)
+if t<=6400000000:
+    print("Running with less than 6gb memory. Working is not garanted")
 
-    ##Generating with stable diffusion
-    device = "cuda"
-    model_path = "CompVis/stable-diffusion-v1-4"
+file1 = open("prompt.txt","r+")
+text = file1.read()
+print(text)
 
-    # Using DDIMScheduler as anexample,this also works with PNDMScheduler
+##Generating with stable diffusion
+device = "cuda"
+model_path = "CompVis/stable-diffusion-v1-4"
+
+    # Using DDIMScheduler as an example,this also works with PNDMScheduler
     # uncomment this line if you want to use it.
 
     #scheduler = PNDMScheduler.from_config(model_path, subfolder="scheduler", use_auth_token=True)
-    scheduler = DDIMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", clip_sample=False, set_alpha_to_one=False)
-    pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
+scheduler = DDIMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", clip_sample=False, set_alpha_to_one=False)
+pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
     model_path,
     scheduler=scheduler,
     revision="fp16", 
     torch_dtype=torch.float16,
     use_auth_token=True
-    ).to(device)
-    if t <= 10500000000:
+).to(device)
+if t <= 11000000000:
+    print("Less then 11gb video memory. Running with attention slicing.")
     pipe.enable_attention_slicing()
         
-    def dummy_checker(images, **kwargs):
+def dummy_checker(images, **kwargs):
     return images, False
-    pipe.safety_checker = dummy_checker
+pipe.safety_checker = dummy_checker
 
-    def preprocess(image):
+def preprocess(image):
     w, h = image.size
     w, h = map(lambda x: x - x % 32, (w, h))  # resize to integer multiple of 32
     image = image.resize((w, h), Image.Resampling.LANCZOS)
@@ -59,39 +60,39 @@ else:
     image = torch.from_numpy(image)
     return 2.*image - 1.
 
-    path_img = []
+path_img = []
 
-    startStrength = 0.86
-    deltaStrength = 0.02
-    endStrength = 0.941
+startStrength = 0.86
+deltaStrength = 0.02
+endStrength = 0.941
 
-    startScale = 7.5
-    deltaScale = 2.5
-    endScale = 15.0
+startScale = 7.5
+deltaScale = 2.5
+endScale = 15.0
 
-    startSeed = 1022
-    endSeed = 1024
+startSeed = 1022
+endSeed = 1024
 
-    directory_in = "./data/input"
-    if not os.path.exists(directory_in):
+directory_in = "./data/input"
+if not os.path.exists(directory_in):
     os.makedirs(directory_in)
-    for root, subdirectories, files in os.walk(directory_in):
+for root, subdirectories, files in os.walk(directory_in):
     for filename in files:
         if filename.endswith(".png"):
            path_img.append(os.path.join(root, filename))
 
-    print("Found " +str(len(path_img)) +" pictures")
-    start_time = time.time()
+print("Found " +str(len(path_img)) +" pictures")
+start_time = time.time()
 
-    counterr=0
-    allwork=0
-    directory="./data/out/" + strftime("%Y-%m-%d_%H-%M-%S", gmtime()) + "/"
-    if not os.path.exists(directory):
+counterr=0
+allwork=0
+directory="./data/out/" + strftime("%Y-%m-%d_%H-%M-%S") + "/"
+if not os.path.exists(directory):
     os.makedirs(directory)
-    with open(directory+"prompt.txt", 'w') as f:
+with open(directory+"prompt.txt", 'w') as f:
     f.write(text)  
-    shutil.copytree(directory_in, directory+"input")
-    for i in path_img:
+shutil.copytree(directory_in, directory+"input")
+for i in path_img:
     print(i)
     if not os.path.exists(directory+str(counterr)):
         os.makedirs(directory+str(counterr))
@@ -129,4 +130,4 @@ else:
         
     counterr+=1
 
-    print("Made " + str(allwork) + " pictures in " + str(time.time()-start_time) + " seconds")
+print("Made " + str(allwork) + " pictures in " + str(time.time()-start_time) + " seconds")