Update app.py

app.py

@@ -1,20 +1,20 @@
 import numpy as np
 import gradio as gr
 import requests
-import random
 import time
 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):
         self.base = base or "https://api.prodia.com/v1"
@@ -34,6 +34,10 @@ class Prodia:
         response = self._post(f"{self.base}/sd/controlnet", params)
         return response.json()
 
+    def upscale(self, params):
+        response = self._post(f"{self.base}/upscale", params)
+        return response.json()
+
     def get_job(self, job_id):
         response = self._get(f"{self.base}/job/{job_id}")
         return response.json()
@@ -42,7 +46,7 @@ class Prodia:
         job_result = job
 
         while job_result['status'] not in ['succeeded', 'failed']:
-            time.sleep(0.25)
+            time.sleep(0.5)
             job_result = self.get_job(job['job'])
 
         return job_result
@@ -76,15 +80,13 @@ class Prodia:
         return response
 
 
-def image_to_base64(image_path):
-    # Open the image with PIL
-    with Image.open(image_path) as image:
-        # Convert the image to bytes
-        buffered = BytesIO()
-        image.save(buffered, format="PNG")  # You can change format to PNG if needed
+def image_to_base64(image):
+    # 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())
+    # Encode the bytes to base64
+    img_str = base64.b64encode(buffered.getvalue())
 
     return img_str.decode('utf-8')  # Convert bytes to string
 
@@ -98,43 +100,6 @@ def remove_id_and_ext(text):
         text = text[:-4]
     return text
 
-def place_lora(current_prompt, lora_name):
-  pattern = r"<lora:" + lora_name + r":.*?>"
-  
-  if re.search(pattern, current_prompt):
-    return re.sub(pattern, "", current_prompt)
-  else:
-    return current_prompt + "<lora:" + lora_name + ":1>"
-
-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 = {}
@@ -199,9 +164,17 @@ def send_to_txt2img(image):
 
         return result
 
+def place_lora(current_prompt, lora_name):
+    pattern = r"<lora:" + lora_name + r":.*?>"
+
+    if re.search(pattern, current_prompt):
+        yield re.sub(pattern, "", current_prompt)
+    else:
+        yield current_prompt + " <lora:" + lora_name + ":1> "
 
 prodia_client = Prodia(api_key=os.getenv("PRODIA_API_KEY"))
 model_list = prodia_client.list_models()
+lora_list = prodia_client.list_loras()
 model_names = {}
 
 for model_name in model_list:
@@ -209,10 +182,11 @@ for model_name in model_list:
     model_names[name_without_ext] = model_name
 
 
-
-
-def flip_text(prompt, negative_prompt, model, steps, sampler, cfg_scale, width, height, seed, batch_size, batch_count, gallery):
-
+def txt2img(prompt, negative_prompt, model, steps, sampler, cfg_scale, width, height, seed, batch_count, gallery):
+    yield {
+        text_button: gr.update(visible=False),
+        stop_btn: gr.update(visible=True),
+    }
     data = {
         "prompt": prompt,
         "negative_prompt": negative_prompt,
@@ -222,41 +196,79 @@ def flip_text(prompt, negative_prompt, model, steps, sampler, cfg_scale, width,
         "cfg_scale": cfg_scale,
         "width": width,
         "height": height,
-        "seed": seed,
-        "upscale": True
+        "seed": seed
     }
 
     total_images = []
-    count_threads = []
+    threads = []
 
-    def generate_one_grid():
-        grid_images = []
-        size_threads = []
+    def generate_one_image():
+        result = prodia_client.generate(data)
+        job = prodia_client.wait(result)
+        total_images.append(job['imageUrl'])
 
-        def generate_one_image():
+    for x in range(batch_count):
+        t = Thread(target=generate_one_image)
+        threads.append(t)
+        t.start()
 
-            result = prodia_client.generate(data)
+    for t in threads:
+        t.join()
+
+    new_images_list = [img['name'] for img in gallery]
+
+    for image in total_images:
+        new_images_list.insert(0, image)
 
-            job = prodia_client.wait(result)
+    if batch_count > 1:
+        results = gr.update(value=total_images, preview=False)
+    else:
+        results = gr.update(value=total_images, preview=True)
+
+    yield {
+        text_button: gr.update(visible=True),
+        stop_btn: gr.update(visible=False),
+        image_output: results,
+        gallery_obj: gr.update(value=new_images_list),
+    }
 
-            grid_images.append(job['imageUrl'])
 
-        for y in range(batch_size):
-            t = Thread(target=generate_one_image)
-            size_threads.append(t)
-            t.start()
+def img2img(input_image, denoising, prompt, negative_prompt, model, steps, sampler, cfg_scale, width, height, seed,
+            batch_count, gallery):
+    if input_image is None:
+        return
+    yield {
+        i2i_text_button: gr.update(visible=False),
+        i2i_stop_btn: gr.update(visible=True),
+    }
+    data = {
+        "imageData": image_to_base64(input_image),
+        "denoising_strength": denoising,
+        "prompt": prompt,
+        "negative_prompt": negative_prompt,
+        "model": model,
+        "steps": steps,
+        "sampler": sampler,
+        "cfg_scale": cfg_scale,
+        "width": width,
+        "height": height,
+        "seed": seed
+    }
 
-        for t in size_threads:
-            t.join()
+    total_images = []
+    threads = []
 
-        total_images.append(create_grid(grid_images))
+    def generate_one_image():
+        result = prodia_client.transform(data)
+        job = prodia_client.wait(result)
+        total_images.append(job['imageUrl'])
 
     for x in range(batch_count):
-        t = Thread(target=generate_one_grid)
-        count_threads.append(t)
+        t = Thread(target=generate_one_image)
+        threads.append(t)
         t.start()
 
-    for t in count_threads:
+    for t in threads:
         t.join()
 
     new_images_list = [img['name'] for img in gallery]
@@ -264,89 +276,274 @@ def flip_text(prompt, negative_prompt, model, steps, sampler, cfg_scale, width,
     for image in total_images:
         new_images_list.insert(0, image)
 
-    return {image_output: total_images, gallery_obj: new_images_list}
+    if batch_count > 1:
+        results = gr.update(value=total_images, preview=False)
+    else:
+        results = gr.update(value=total_images, preview=True)
 
+    yield {
+        i2i_text_button: gr.update(visible=True),
+        i2i_stop_btn: gr.update(visible=False),
+        i2i_image_output: results,
+        gallery_obj: gr.update(value=new_images_list),
+    }
+
+def upscale_fn(image, scale):
+    if image is None:
+        return
+    yield {
+        upscale_btn: gr.update(visible=False),
+        upscale_stop: gr.update(visible=True),
+    }
+    job = prodia_client.upscale({
+        'imageData': image_to_base64(image),
+        'resize': scale
+    })
+
+    result = prodia_client.wait(job)
+    yield {
+        upscale_output: result['imageUrl'],
+        upscale_btn: gr.update(visible=True),
+        upscale_stop: gr.update(visible=False)
+    }
+
+def stop_upscale():
+    return {
+        upscale_btn: gr.update(visible=True),
+        upscale_stop: gr.update(visible=False)
+    }
+
+def stop_t2i():
+    return {
+        text_button: gr.update(visible=True),
+        stop_btn: gr.update(visible=False)
+    }
+
+def stop_i2i():
+    return {
+        i2i_text_button: gr.update(visible=True),
+        i2i_stop_btn: gr.update(visible=False)
+    }
+
+
+
+samplers = [
+    "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",
+]
 
 css = """
-#model_dropdown {
-    width: 50%;
+:root, .dark{
+    --checkbox-label-gap: 0.25em 0.1em;
+    --section-header-text-size: 12pt;
+    --block-background-fill: transparent;
+}
+.block.padded:not(.gradio-accordion) {
+    padding: 0 !important;
+}
+div.gradio-container{
+    max-width: unset !important;
+}
+.compact{
+    background: transparent !important;
+    padding: 0 !important;
+}
+div.form{
+    border-width: 0;
+    box-shadow: none;
+    background: transparent;
+    overflow: visible;
+    gap: 0.5em;
+}
+.block.gradio-dropdown,
+.block.gradio-slider,
+.block.gradio-checkbox,
+.block.gradio-textbox,
+.block.gradio-radio,
+.block.gradio-checkboxgroup,
+.block.gradio-number,
+.block.gradio-colorpicker {
+    border-width: 0 !important;
+    box-shadow: none !important;
+}
+.gradio-dropdown label span:not(.has-info),
+.gradio-textbox label span:not(.has-info),
+.gradio-number label span:not(.has-info)
+{
+    margin-bottom: 0;
+}
+.gradio-dropdown ul.options{
+    z-index: 3000;
+    min-width: fit-content;
+    max-width: inherit;
+    white-space: nowrap;
+}
+.gradio-dropdown ul.options li.item {
+    padding: 0.05em 0;
+}
+.gradio-dropdown ul.options li.item.selected {
+    background-color: var(--neutral-100);
+}
+.dark .gradio-dropdown ul.options li.item.selected {
+    background-color: var(--neutral-900);
+}
+.gradio-dropdown div.wrap.wrap.wrap.wrap{
+    box-shadow: 0 1px 2px 0 rgba(0, 0, 0, 0.05);
+}
+.gradio-dropdown:not(.multiselect) .wrap-inner.wrap-inner.wrap-inner{
+    flex-wrap: unset;
+}
+.gradio-dropdown .single-select{
+    white-space: nowrap;
+    overflow: hidden;
+}
+.gradio-dropdown .token-remove.remove-all.remove-all{
+    display: none;
+}
+.gradio-dropdown.multiselect .token-remove.remove-all.remove-all{
+    display: flex;
+}
+.gradio-slider input[type="number"]{
+    width: 6em;
+}
+.block.gradio-checkbox {
+    margin: 0.75em 1.5em 0 0;
+}
+.gradio-html div.wrap{
+    height: 100%;
+}
+div.gradio-html.min{
+    min-height: 0;
+}
+#model_dd {
+    width: 16%;
 }
 """
 
 with gr.Blocks(css=css) as demo:
-    with gr.Row(elem_id="model_dropdown"):
-        with gr.Column(scale=2):
-            model = gr.Dropdown(interactive=True, value="absolutereality_v181.safetensors [3d9d4d2b]",
-                                label="Stable Diffusion Checkpoint", choices=prodia_client.list_models(), container=False)
-
-        with gr.Column(scale=6):
-            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)")
+    model = gr.Dropdown(interactive=True, value="absolutereality_v181.safetensors [3d9d4d2b]", show_label=True,
+                                label="Stable Diffusion Checkpoint", choices=prodia_client.list_models(), elem_id="model_dd")
 
     with gr.Tabs() as tabs:
-        with gr.Tab("txt2img", id='t2i', container=False):
+        with gr.Tab("txt2img", id='t2i'):
             with gr.Row():
-                with gr.Column(scale=4, min_width=600):
+                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, container=False)
-                    negative_prompt = gr.Textbox(placeholder="Negative Prompt", show_label=False, lines=3, container=False,
+                                        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.Row(scale=6):
+                with gr.Row():
                     text_button = gr.Button("Generate", variant='primary', elem_id="generate")
-                    stop_btn = gr.Button("Cancel", elem_id="cancel")
+                    stop_btn = gr.Button("Cancel", variant="stop", elem_id="generate", visible=False)
 
             with gr.Row():
-                with gr.Column(scale=3):
+                with gr.Column():
                     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="DPM++ 2M Karras", show_label=True, label="Sampling Method",
+                                                      choices=samplers)
 
                             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):
+                            with gr.Column(scale=8):
                                 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", minimum=1, maximum=9, value=1, step=1)
-                                batch_count = gr.Slider(label="Batch Count", minimum=1, maximum=100, value=1, step=1)
+                                batch_size = gr.Slider(label="Batch Size", maximum=1, value=1)
+                                batch_count = gr.Slider(label="Batch Count", minimum=1, maximum=4, 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.Tab("Lora"):
-                        loralist = prodia_client.list_loras()
                         with gr.Row():
-                            for lora in loralist:
+                            for lora in lora_list:
                                 lora_btn = gr.Button(lora, size="sm")
                                 lora_btn.click(place_lora, inputs=[prompt, lora_btn], outputs=prompt)
 
-                with gr.Column(scale=2):
-                    image_output = gr.Gallery(value=["https://images.prodia.xyz/8ede1a7c-c0ee-4ded-987d-6ffed35fc477.png"], preview=True)
+                with gr.Column():
+                    image_output = gr.Gallery(columns=3,
+                        value=["https://images.prodia.xyz/8ede1a7c-c0ee-4ded-987d-6ffed35fc477.png"])
+
+        with gr.Tab("img2img", id='i2i'):
+            with gr.Row():
+                with gr.Column(scale=6, min_width=600):
+                    i2i_prompt = gr.Textbox("space warrior, beautiful, female, ultrarealistic, soft lighting, 8k",
+                                            placeholder="Prompt", show_label=False, lines=3)
+                    i2i_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.Row():
+                    i2i_text_button = gr.Button("Generate", variant='primary', elem_id="generate")
+                    i2i_stop_btn = gr.Button("Cancel", variant="stop", elem_id="generate", visible=False)
+
+            with gr.Row():
+                with gr.Column(scale=1):
+                    with gr.Tab("Generation"):
+                        i2i_image_input = gr.Image(type="pil")
+
+                        with gr.Row():
+                            with gr.Column(scale=1):
+                                i2i_sampler = gr.Dropdown(value="DPM++ 2M Karras", show_label=True,
+                                                          label="Sampling Method", choices=samplers)
+
+                            with gr.Column(scale=1):
+                                i2i_steps = gr.Slider(label="Sampling Steps", minimum=1, maximum=30, value=25, step=1)
+
+                        with gr.Row():
+                            with gr.Column(scale=6):
+                                i2i_width = gr.Slider(label="Width", maximum=1024, value=512, step=8)
+                                i2i_height = gr.Slider(label="Height", maximum=1024, value=512, step=8)
+
+                            with gr.Column(scale=1):
+                                i2i_batch_size = gr.Slider(label="Batch Size", maximum=1, value=1)
+                                i2i_batch_count = gr.Slider(label="Batch Count", minimum=1, maximum=4, value=1, step=1)
+
+                        i2i_cfg_scale = gr.Slider(label="CFG Scale", minimum=1, maximum=20, value=7, step=1)
+                        i2i_denoising = gr.Slider(label="Denoising Strength", minimum=0, maximum=1, value=0.7, step=0.1)
+                        i2i_seed = gr.Number(label="Seed", value=-1)
+
+                    with gr.Tab("Lora"):
+                        with gr.Row():
+                            for lora in lora_list:
+                                lora_btn = gr.Button(lora, size="sm")
+                                lora_btn.click(place_lora, inputs=[i2i_prompt, lora_btn], outputs=i2i_prompt)
+
+                with gr.Column(scale=1):
+                    i2i_image_output = gr.Gallery(columns=3,
+                        value=["https://images.prodia.xyz/8ede1a7c-c0ee-4ded-987d-6ffed35fc477.png"])
+
+        with gr.Tab("Extras"):
+            with gr.Row():
+                with gr.Tab("Single Image"):
+                    with gr.Column():
+                        upscale_image_input = gr.Image(type="pil")
+                        upscale_btn = gr.Button("Generate", variant="primary")
+                        upscale_stop = gr.Button("Stop", variant="stop", visible=False)
+                        with gr.Tab("Scale by"):
+                            scale_by = gr.Radio([2, 4], value=2, label="Resize")
+
+                upscale_output = gr.Image()
 
         with gr.Tab("PNG Info"):
             def plaintext_to_html(text, classname=None):
@@ -373,6 +570,7 @@ with gr.Blocks(css=css) as demo:
 
                 return info
 
+
             with gr.Row():
                 with gr.Column():
                     image_input = gr.Image(type="pil")
@@ -382,16 +580,26 @@ with gr.Blocks(css=css) as demo:
                     send_to_txt2img_btn = gr.Button("Send to txt2img")
 
         with gr.Tab("Gallery"):
-            gallery_obj = gr.Gallery(height=1000, columns=5)
+            gallery_obj = gr.Gallery(height=1000, columns=6)
+
+        t2i_event = text_button.click(txt2img,
+                                      inputs=[prompt, negative_prompt, model, steps, sampler, cfg_scale, width, height,
+                                              seed, batch_count, gallery_obj], outputs=[image_output, gallery_obj, text_button, stop_btn])
+        stop_btn.click(fn=stop_t2i, outputs=[text_button, stop_btn], cancels=[t2i_event])
 
-        generation_event = 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])
-        stop_btn.click(fn=None, outputs=None, cancels=[generation_event])
         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])
+                                  outputs=[tabs, prompt, negative_prompt, steps, seed, model, sampler, width, height,
+                                           cfg_scale])
+
+        i2i_event = i2i_text_button.click(img2img,
+                                          inputs=[i2i_image_input, i2i_denoising, i2i_prompt, i2i_negative_prompt,
+                                                  model, i2i_steps, i2i_sampler, i2i_cfg_scale, i2i_width, i2i_height,
+                                                  i2i_seed, i2i_batch_count, gallery_obj],
+                                          outputs=[i2i_image_output, gallery_obj, i2i_text_button, i2i_stop_btn])
+        i2i_stop_btn.click(fn=stop_i2i, outputs=[i2i_text_button, i2i_stop_btn], cancels=[i2i_event])
+
+        upscale_event = upscale_btn.click(fn=upscale_fn, inputs=[upscale_image_input, scale_by], outputs=[upscale_output, upscale_btn, upscale_stop])
+        upscale_stop.click(fn=stop_upscale, outputs=[upscale_btn, upscale_stop], cancels=[upscale_event])
 
-demo.queue(concurrency_count=32)
-demo.launch()
+demo.queue(concurrency_count=64, max_size=80, api_open=False).launch(max_threads=256)