Update app.py

app.py

@@ -6,11 +6,16 @@ import json
 import base64
 import os
 from io import BytesIO
+import math
 import PIL
+from PIL import Image
 from PIL.ExifTags import TAGS
 import html
 import re
+from threading import Thread
 
+from dotenv import load_dotenv
+load_dotenv()
 
 class Prodia:
     def __init__(self, api_key, base=None):
@@ -18,19 +23,19 @@ class Prodia:
         self.headers = {
             "X-Prodia-Key": api_key
         }
-    
+
     def generate(self, params):
         response = self._post(f"{self.base}/sd/generate", params)
         return response.json()
-    
+
     def transform(self, params):
         response = self._post(f"{self.base}/sd/transform", params)
         return response.json()
-    
+
     def controlnet(self, params):
         response = self._post(f"{self.base}/sd/controlnet", params)
         return response.json()
-    
+
     def get_job(self, job_id):
         response = self._get(f"{self.base}/job/{job_id}")
         return response.json()
@@ -75,12 +80,13 @@ def image_to_base64(image_path):
         # Convert the image to bytes
         buffered = BytesIO()
         image.save(buffered, format="PNG")  # You can change format to PNG if needed
-        
+
         # Encode the bytes to base64
         img_str = base64.b64encode(buffered.getvalue())
 
     return img_str.decode('utf-8')  # Convert bytes to string
 
+
 def remove_id_and_ext(text):
     text = re.sub(r'\[.*\]$', '', text)
     extension = text[-12:].strip()
@@ -90,6 +96,37 @@ def remove_id_and_ext(text):
         text = text[:-4]
     return text
 
+
+def create_grid(image_urls):
+    # Download first image to get size
+    response = requests.get(image_urls[0])
+    img_data = response.content
+    img = Image.open(BytesIO(img_data))
+    w, h = img.size
+
+    # Calculate rows and cols
+    num_images = len(image_urls)
+    num_cols = min(num_images, 3)
+    num_rows = math.ceil(num_images / num_cols)
+
+    # Create new rgba image
+    grid_w = num_cols * w
+    grid_h = num_rows * h
+    grid = Image.new('RGBA', (grid_w, grid_h), (0, 0, 0, 0))
+
+    # Download images and paste into grid
+    for index, img_url in enumerate(image_urls):
+        response = requests.get(img_url)
+        img_data = response.content
+        img = Image.open(BytesIO(img_data))
+
+        row = index // num_cols
+        col = index % num_cols
+        grid.paste(img, (col * w, row * h))
+
+    # Save image
+    return grid
+
 def get_data(text):
     results = {}
     patterns = {
@@ -97,11 +134,11 @@ def get_data(text):
         'negative_prompt': r'Negative prompt: (.*)',
         'steps': r'Steps: (\d+),',
         'seed': r'Seed: (\d+),',
-        'sampler': r'Sampler:\s*([^\s,]+(?:\s+[^\s,]+)*)', 
+        'sampler': r'Sampler:\s*([^\s,]+(?:\s+[^\s,]+)*)',
         'model': r'Model:\s*([^\s,]+)',
         'cfg_scale': r'CFG scale:\s*([\d\.]+)',
         'size': r'Size:\s*([0-9]+x[0-9]+)'
-        }
+    }
     for key in ['prompt', 'negative_prompt', 'steps', 'seed', 'sampler', 'model', 'cfg_scale', 'size']:
         match = re.search(patterns[key], text)
         if match:
@@ -117,23 +154,24 @@ def get_data(text):
         results['h'] = None
     return results
 
+
 def send_to_txt2img(image):
-    
     result = {tabs: gr.Tabs.update(selected="t2i")}
 
     try:
         text = image.info['parameters']
         data = get_data(text)
         result[prompt] = gr.update(value=data['prompt'])
-        result[negative_prompt] = gr.update(value=data['negative_prompt']) if data['negative_prompt'] is not None else gr.update()
+        result[negative_prompt] = gr.update(value=data['negative_prompt']) if data[
+                                                                                  'negative_prompt'] is not None else gr.update()
         result[steps] = gr.update(value=int(data['steps'])) if data['steps'] is not None else gr.update()
         result[seed] = gr.update(value=int(data['seed'])) if data['seed'] is not None else gr.update()
         result[cfg_scale] = gr.update(value=float(data['cfg_scale'])) if data['cfg_scale'] is not None else gr.update()
         result[width] = gr.update(value=int(data['w'])) if data['w'] is not None else gr.update()
         result[height] = gr.update(value=int(data['h'])) if data['h'] is not None else gr.update()
         result[sampler] = gr.update(value=data['sampler']) if data['sampler'] is not None else gr.update()
-        if model in model_names:
-            result[model] = gr.update(value=model_names[model])
+        if data['model'] in model_names:
+            result[model] = gr.update(value=model_names[data['model']])
         else:
             result[model] = gr.update()
         return result
@@ -153,7 +191,6 @@ def send_to_txt2img(image):
         return result
 
 
-
 prodia_client = Prodia(api_key=os.getenv("PRODIA_API_KEY"))
 model_list = prodia_client.list_models()
 model_names = {}
@@ -162,8 +199,12 @@ for model_name in model_list:
     name_without_ext = remove_id_and_ext(model_name)
     model_names[name_without_ext] = model_name
 
-def flip_text(prompt, negative_prompt, model, steps, sampler, cfg_scale, width, height, seed):
-    result = prodia_client.generate({
+
+
+
+def flip_text(prompt, negative_prompt, model, steps, sampler, cfg_scale, width, height, seed, batch_size, batch_count, gallery):
+
+    data = {
         "prompt": prompt,
         "negative_prompt": negative_prompt,
         "model": model,
@@ -173,11 +214,47 @@ def flip_text(prompt, negative_prompt, model, steps, sampler, cfg_scale, width,
         "width": width,
         "height": height,
         "seed": seed
-    })
+    }
+
+    total_images = []
+    count_threads = []
+
+    def generate_one_grid():
+        grid_images = []
+        size_threads = []
+
+        def generate_one_image():
+
+            result = prodia_client.generate(data)
+
+            job = prodia_client.wait(result)
+
+            grid_images.append(job['imageUrl'])
 
-    job = prodia_client.wait(result)
+        for y in range(batch_size):
+            t = Thread(target=generate_one_image)
+            size_threads.append(t)
+            t.start()
 
-    return job["imageUrl"]
+        for t in size_threads:
+            t.join()
+
+        total_images.append(create_grid(grid_images))
+
+    for x in range(batch_count):
+        t = Thread(target=generate_one_grid)
+        count_threads.append(t)
+        t.start()
+
+    for t in count_threads:
+        t.join()
+
+    new_images_list = [img['name'] for img in gallery]
+
+    for image in total_images:
+        new_images_list.insert(0, image)
+
+    return {image_output: total_images, gallery_obj: new_images_list}
 
 
 css = """
@@ -187,82 +264,82 @@ css = """
 """
 
 with gr.Blocks(css=css) as demo:
-
-    
     with gr.Row():
         with gr.Column(scale=6):
-            model = gr.Dropdown(interactive=True,value="absolutereality_v181.safetensors [3d9d4d2b]", show_label=True, label="Stable Diffusion Checkpoint", choices=prodia_client.list_models())
-  
+            model = gr.Dropdown(interactive=True, value="absolutereality_v181.safetensors [3d9d4d2b]", show_label=True,
+                                label="Stable Diffusion Checkpoint", choices=prodia_client.list_models())
+
         with gr.Column(scale=1):
-            gr.Markdown(elem_id="powered-by-prodia", value="AUTOMATIC1111 Stable Diffusion Web UI.<br>Powered by [Prodia](https://prodia.com).<br> For more features and faster gen times check out our [API Docs](https://docs.prodia.com/reference/getting-started-guide)")
+            gr.Markdown(elem_id="powered-by-prodia",
+                        value="AUTOMATIC1111 Stable Diffusion Web UI.<br>Powered by [Prodia](https://prodia.com).<br> For more features and faster gen times check out our [API Docs](https://docs.prodia.com/reference/getting-started-guide)")
 
     with gr.Tabs() as tabs:
         with gr.Tab("txt2img", id='t2i'):
             with gr.Row():
                 with gr.Column(scale=6, min_width=600):
-                    prompt = gr.Textbox("space warrior, beautiful, female, ultrarealistic, soft lighting, 8k", placeholder="Prompt", show_label=False, lines=3)
-                    negative_prompt = gr.Textbox(placeholder="Negative Prompt", show_label=False, lines=3, value="3d, cartoon, anime, (deformed eyes, nose, ears, nose), bad anatomy, ugly")
+                    prompt = gr.Textbox("space warrior, beautiful, female, ultrarealistic, soft lighting, 8k",
+                                        placeholder="Prompt", show_label=False, lines=3)
+                    negative_prompt = gr.Textbox(placeholder="Negative Prompt", show_label=False, lines=3,
+                                                 value="3d, cartoon, anime, (deformed eyes, nose, ears, nose), bad anatomy, ugly")
                 with gr.Column():
                     text_button = gr.Button("Generate", variant='primary', elem_id="generate")
-                    
+
             with gr.Row():
                 with gr.Column(scale=3):
                     with gr.Tab("Generation"):
                         with gr.Row():
                             with gr.Column(scale=1):
-                                sampler = gr.Dropdown(value="Euler a", show_label=True, label="Sampling Method", choices=[
-                                    "Euler",
-                                    "Euler a",
-                                    "LMS",
-                                    "Heun",
-                                    "DPM2",
-                                    "DPM2 a",
-                                    "DPM++ 2S a",
-                                    "DPM++ 2M",
-                                    "DPM++ SDE",
-                                    "DPM fast",
-                                    "DPM adaptive",
-                                    "LMS Karras",
-                                    "DPM2 Karras",
-                                    "DPM2 a Karras",
-                                    "DPM++ 2S a Karras",
-                                    "DPM++ 2M Karras",
-                                    "DPM++ SDE Karras",
-                                    "DDIM",
-                                    "PLMS",
-                                ])
-                                
+                                sampler = gr.Dropdown(value="Euler a", show_label=True, label="Sampling Method",
+                                                      choices=[
+                                                          "Euler",
+                                                          "Euler a",
+                                                          "LMS",
+                                                          "Heun",
+                                                          "DPM2",
+                                                          "DPM2 a",
+                                                          "DPM++ 2S a",
+                                                          "DPM++ 2M",
+                                                          "DPM++ SDE",
+                                                          "DPM fast",
+                                                          "DPM adaptive",
+                                                          "LMS Karras",
+                                                          "DPM2 Karras",
+                                                          "DPM2 a Karras",
+                                                          "DPM++ 2S a Karras",
+                                                          "DPM++ 2M Karras",
+                                                          "DPM++ SDE Karras",
+                                                          "DDIM",
+                                                          "PLMS",
+                                                      ])
+
                             with gr.Column(scale=1):
                                 steps = gr.Slider(label="Sampling Steps", minimum=1, maximum=30, value=25, step=1)
-    
+
                         with gr.Row():
                             with gr.Column(scale=1):
                                 width = gr.Slider(label="Width", maximum=1024, value=512, step=8)
                                 height = gr.Slider(label="Height", maximum=1024, value=512, step=8)
-                            
+
                             with gr.Column(scale=1):
-                                batch_size = gr.Slider(label="Batch Size", maximum=1, value=1)
-                                batch_count = gr.Slider(label="Batch Count", maximum=1, value=1)
-    
+                                batch_size = gr.Slider(label="Batch Size", minimum=1, maximum=9, value=1, step=1)
+                                batch_count = gr.Slider(label="Batch Count", minimum=1, maximum=100, value=1, step=1)
+
                         cfg_scale = gr.Slider(label="CFG Scale", minimum=1, maximum=20, value=7, step=1)
                         seed = gr.Number(label="Seed", value=-1)
-    
-                    
+
                 with gr.Column(scale=2):
-                    image_output = gr.Image(value="https://images.prodia.xyz/8ede1a7c-c0ee-4ded-987d-6ffed35fc477.png")
-    
-            text_button.click(flip_text, inputs=[prompt, negative_prompt, model, steps, sampler, cfg_scale, width, height, seed], outputs=image_output)
-        
+                    image_output = gr.Gallery(value=["https://images.prodia.xyz/8ede1a7c-c0ee-4ded-987d-6ffed35fc477.png"], preview=True)
+
         with gr.Tab("PNG Info"):
             def plaintext_to_html(text, classname=None):
                 content = "<br>\n".join(html.escape(x) for x in text.split('\n'))
-    
+
                 return f"<p class='{classname}'>{content}</p>" if classname else f"<p>{content}</p>"
-    
-    
+
+
             def get_exif_data(image):
                 items = image.info
-    
+
                 info = ''
                 for key, text in items.items():
                     info += f"""
@@ -270,25 +347,32 @@ with gr.Blocks(css=css) as demo:
                     <p><b>{plaintext_to_html(str(key))}</b></p>
                     <p>{plaintext_to_html(str(text))}</p>
                     </div>
-                    """.strip()+"\n"
-    
+                    """.strip() + "\n"
+
                 if len(info) == 0:
                     message = "Nothing found in the image."
                     info = f"<div><p>{message}<p></div>"
-    
+
                 return info
-    
+
             with gr.Row():
                 with gr.Column():
                     image_input = gr.Image(type="pil")
-                    
+
                 with gr.Column():
                     exif_output = gr.HTML(label="EXIF Data")
                     send_to_txt2img_btn = gr.Button("Send to txt2img")
-    
-            image_input.upload(get_exif_data, inputs=[image_input], outputs=exif_output)
-            send_to_txt2img_btn.click(send_to_txt2img, inputs=[image_input], outputs=[tabs, prompt, negative_prompt, steps, seed,
-                                                                                          model, sampler, width, height, cfg_scale])
+
+        with gr.Tab("Gallery"):
+            gallery_obj = gr.Gallery(height=1000, columns=5)
+
+        text_button.click(flip_text,
+                          inputs=[prompt, negative_prompt, model, steps, sampler, cfg_scale, width, height, seed, batch_size, batch_count,
+                                  gallery_obj], outputs=[image_output, gallery_obj])
+        image_input.upload(get_exif_data, inputs=[image_input], outputs=exif_output)
+        send_to_txt2img_btn.click(send_to_txt2img, inputs=[image_input],
+                                  outputs=[tabs, prompt, negative_prompt, steps, seed,
+                                           model, sampler, width, height, cfg_scale])
 
 demo.queue(concurrency_count=32)
-demo.launch()
+demo.launch()