batch size option added

Author: divamgupta

backends/stable_diffusion_tf/diffusionbee_backend.py

@@ -7,6 +7,7 @@
 from downloader import ProgressBarDownloader
 import sys
 import copy
+import math
 
 class Unbuffered(object):
     def __init__(self, stream):
@@ -29,8 +30,12 @@ def __getattr__(self, attr):
 
 
 def process_opt(d, generator):
+
+    batch_size = int(d['batch_size'])
+    n_imgs = math.ceil(d['num_imgs'] / batch_size)
+
 
-    for i in range(d['num_imgs']):
+    for i in range(n_imgs):
         if 'seed' in d:
             seed = d['seed']
         else:
@@ -41,15 +46,17 @@ def process_opt(d, generator):
             num_steps=d['ddim_steps'],
             unconditional_guidance_scale=d['scale'],
             temperature=1,
-            batch_size=1,
+            batch_size=batch_size,
             seed=seed,
             img_id=i,
         )
         if img is None:
             return
-        fpath = "/tmp/%d.png"%(random.randint(0 ,100000000))
-        Image.fromarray(img[0]).save(fpath)
-        print("sdbk nwim %s"%(fpath) )
+        
+        for i in range(len(img)):
+            fpath = "/tmp/%d.png"%(random.randint(0 ,100000000))
+            Image.fromarray(img[i]).save(fpath)
+            print("sdbk nwim %s"%(fpath) )
 
 
 def main():
@@ -88,7 +95,7 @@ def main():
     generator.diffusion_model.load_weights(p1)  
     generator.decoder.load_weights(p3) 
 
-    default_d = { "W" : 512 , "H" : 512, "num_imgs":1 , "ddim_steps" : 25 , "scale" : 7.5 }
+    default_d = { "W" : 512 , "H" : 512, "num_imgs":1 , "ddim_steps" : 25 , "scale" : 7.5, "batch_size":1 }
 
 
     print("sdbk mdld")

electron_app/src/components/ImgGenerate.vue

@@ -72,6 +72,16 @@
                                 ></b-form-select>
                                 </b-form-group>
 
+                                <b-form-group inline label=""  style="margin-bottom: 6px;">
+                                <label class="mr-sm-2" style="margin-right: 8px ;" for="inline-form-custom-select-pref">Batch size: </label>
+                                <b-form-select
+                                    style="border-color:rgba(0,0,0,0.1)"
+                                    v-model="batch_size"
+                                    :options="[1, 2, 3, 4, 5, 6, 7, 8]"
+                                    required
+                                ></b-form-select>
+                                </b-form-group>
+
                                 <b-form-group inline  label="" style="margin-bottom: 6px;" >
                                 <label class="mr-sm-2" style="margin-right: 8px ;" for="inline-form-custom-select-pref">Guidance Scale: </label>
                                 <b-form-select
@@ -82,6 +92,9 @@
                                 ></b-form-select>
                                 </b-form-group>
 
+                                
+
+
                                 <b-form-group inline  label="" style="margin-bottom: 6px;" >
                                 <label class="mr-sm-2" style="margin-right: 8px ;" for="inline-form-custom-select-pref">Seed: </label>
 
@@ -211,6 +224,7 @@ export default {
             seed : ""  , 
             prompt : "",
             num_imgs : 1,
+            batch_size : 1 , 
             generated_images : [],
             backend_error : "",
             done_percentage : -1,
@@ -229,6 +243,7 @@ export default {
                 scale : this.guidence_scale , 
                 ddim_steps : this.dif_steps, 
                 num_imgs : this.num_imgs , 
+                batch_size : this.batch_size , 
 
             }
             let that = this;