Update app.py

app.py

@@ -74,13 +74,25 @@ class Prodia:
         response = self._post(f"{self.base}/sd/generate", params)
         return response.json()
 
+    def xl_generate(self, params):
+        response = self._post(f"{self.base}/sdxl/generate", params)
+        return response.json()
+
     def transform(self, params):
         response = self._post(f"{self.base}/sd/transform", params)
         return response.json()
+    
+    def xl_transform(self, params):
+        response = self._post(f"{self.base}/sdxl/transform", params)
+        return response.json()
 
     def controlnet(self, params):
         response = self._post(f"{self.base}/sd/controlnet", params)
         return response.json()
+    
+    def xl_controlnet(self, params):
+        response = self._post(f"{self.base}/sdxl/controlnet", params)
+        return response.json()
 
     def get_job(self, job_id):
         response = self._get(f"{self.base}/job/{job_id}")
@@ -98,10 +110,18 @@ class Prodia:
     def list_models(self):
         response = self._get(f"{self.base}/sd/models")
         return response.json()
+    
+    def xl_list_models(self):
+        response = self._get(f"{self.base}/sdxl/models")
+        return response.json()
 
     def list_samplers(self):
         response = self._get(f"{self.base}/sd/samplers")
         return response.json()
+    
+    def xl_list_samplers(self):
+        response = self._get(f"{self.base}/sdxl/samplers")
+        return response.json()
 
     def _post(self, url, params):
         headers = {
@@ -207,11 +227,17 @@ def send_to_txt2img(image):
 prodia_client = Prodia(api_key=os.getenv("PRODIA_API_KEY"))
 model_list = prodia_client.list_models()
 model_names = {}
+xl_model_list = prodia_client.xl_list_models()
+xl_model_names = {}
 
 for model_name in model_list:
     name_without_ext = remove_id_and_ext(model_name)
     model_names[name_without_ext] = model_name
 
+for xl_model_name in xl_model_list:
+    xl_name_without_ext = remove_id_and_ext(xl_model_name)
+    xl_model_names[xl_name_without_ext] = xl_model_name
+
 
 def txt2img(prompt, negative_prompt, model, steps, sampler, cfg_scale, width, height, upscale, seed, progress=gr.Progress()):
     progress(0, desc="Starting")
@@ -234,6 +260,25 @@ def txt2img(prompt, negative_prompt, model, steps, sampler, cfg_scale, width, he
     progress(0.99, desc="Sending image")
     return [job["imageUrl"]], job["imageUrl"]
 
+def xl_txt2img(prompt, negative_prompt, model, steps, sampler, cfg_scale, width, height, seed, progress=gr.Progress()):
+    progress(0, desc="Starting")
+    time.sleep(2.5)
+    progress(0.25, desc="Generating")
+    result = prodia_client.generate({
+        "prompt": prompt,
+        "negative_prompt": negative_prompt,
+        "model": model,
+        "steps": steps,
+        "sampler": sampler,
+        "cfg_scale": cfg_scale,
+        "width": width,
+        "height": height,
+        "seed": seed
+    })
+    progress(0.75, desc="Opening image")
+    job = prodia_client.wait(result)
+    progress(0.99, desc="Sending image")
+    return [job["imageUrl"]], job["imageUrl"]
 
 def img2img(input_image, denoising, prompt, negative_prompt, model, steps, sampler, cfg_scale, width, height, upscale, seed, progress=gr.Progress()):
     progress(0, desc="Starting")
@@ -262,6 +307,36 @@ def img2img(input_image, denoising, prompt, negative_prompt, model, steps, sampl
     time.sleep(0.5)
     return [job["imageUrl"]], job["imageUrl"]
 
+def xl_img2img(input_image, denoising, prompt, negative_prompt, model, steps, sampler, cfg_scale, width, height, seed, progress=gr.Progress()):
+    progress(0, desc="Starting")
+    time.sleep(1.5)
+    progress(0.10, desc="Uploading input image")
+    time.sleep(1.5)
+    progress(0.25, desc="Generating")
+    result = prodia_client.transform({
+        "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
+    })
+    progress(0.75, desc="Opening image")
+
+    job = prodia_client.wait(result)
+    progress(0.99, desc="Sending image")
+    time.sleep(0.5)
+    return [job["imageUrl"]], job["imageUrl"]
+
+
+
+
+
 
 css = """
 #generate {
@@ -422,6 +497,99 @@ with gr.Blocks(css=css) as demo:
         send_to_png.click(send_to_img2img_def, inputs=past_url, outputs=image_input)
         send_to_img2img_png.click(send_to_img2img_def, inputs=past_url, outputs=i2i_image_input)
         send_to_png_i2i.click(send_to_img2img_def, inputs=i2i_past_url, outputs=image_input)
+        with gr.Tab("XL Inference"):
+            with gr.Row():
+                with gr.Column(scale=6):
+                    xl_model = gr.Dropdown(interactive=True, show_label=True,
+                                        label="Stable Diffusion XL Checkpoint", choices=prodia_client.xl_list_models())
+
+            with gr.Tab("txt2img"):
+                with gr.Row():
+                    with gr.Column(scale=6, min_width=600):
+                        xl_prompt = gr.Textbox(placeholder="Prompt", show_label=False, lines=3, scale=3)
+                        xl_negative_prompt = gr.Textbox(placeholder="Negative Prompt", show_label=False, lines=3, scale=1,
+                                                         value="(deformed, distorted, disfigured:1.3), poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, (mutated hands and fingers:1.4), disconnected limbs, mutation, mutated, ugly, disgusting, blurry, amputation")
+                    with gr.Column():
+                        xl_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):
+                                    xl_sampler = gr.Dropdown(value="DPM++ 2M Karras", show_label=True, label="Sampling Method",
+                                                            choices=prodia_client.xl_list_samplers())
+    
+                                with gr.Column(scale=1):
+                                    xl_steps = gr.Slider(label="Sampling Steps", minimum=1, maximum=25, value=20, step=1)
+
+                            with gr.Row():
+                                with gr.Column(scale=1):
+                                    xl_width = gr.Slider(label="Width", maximum=1024, value=512, step=8)
+                                    xl_height = gr.Slider(label="Height", maximum=1024, value=512, step=8)
+
+                                with gr.Column(scale=1):
+                                    xl_batch_size = gr.Slider(label="Batch Size", maximum=1, value=1)
+                                    xl_batch_count = gr.Slider(label="Batch Count", maximum=1, value=1)
+
+                            xl_cfg_scale = gr.Slider(label="CFG Scale", minimum=0.1, maximum=20, value=8, step=0.1)
+                            xl_seed = gr.Number(label="Seed", value=-1)
+
+                    with gr.Column(scale=2):
+                        xl_image_output = gr.Gallery(show_label=False, rows=2, preview=True)
+                        xl_send_to_img2img = gr.Button(value="Send OUTPUT IMAGE to img2img")
+                        xl_send_to_png = gr.Button(value="Send OUTPUT IMAGE to PNG Info")
+                        xl_past_url = gr.Textbox(visible=False, interactive=False)
+
+                xl_text_button.click(xl_txt2img, inputs=[xl_prompt, xl_negative_prompt, xl_model, xl_steps, xl_sampler, xl_cfg_scale, xl_width, xl_height,
+                                                      xl_seed], outputs=[xl_image_output, xl_past_url], concurrency_limit=64)
+
+            with gr.Tab("img2img"):
+                with gr.Row():
+                    with gr.Column(scale=6, min_width=600):
+                        xl_i2i_prompt = gr.Textbox(placeholder="Prompt", show_label=False, lines=3, scale=3)
+                        xl_i2i_negative_prompt = gr.Textbox(placeholder="Negative Prompt", show_label=False, lines=3, scale=1,
+                                                     value="(deformed, distorted, disfigured:1.3), poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, (mutated hands and fingers:1.4), disconnected limbs, mutation, mutated, ugly, disgusting, blurry, amputation")
+                    with gr.Column():
+                        xl_i2i_text_button = gr.Button("Generate", variant='primary', elem_id="generate")
+
+                with gr.Row():
+                    with gr.Column(scale=3):
+                        with gr.Tab("Generation"):
+                            xl_i2i_image_input = gr.Image(type="pil", interactive=True)
+
+                            with gr.Row():
+                                with gr.Column(scale=1):
+                                    xl_i2i_sampler = gr.Dropdown(value="DPM++ 2M Karras", show_label=True, label="Sampling Method",
+                                                              choices=prodia_client.xl_list_samplers())
+
+                                with gr.Column(scale=1):
+                                    xl_i2i_steps = gr.Slider(label="Sampling Steps", minimum=1, maximum=25, value=20, step=1)
+                            with gr.Row():
+                                with gr.Column(scale=1):
+                                    xl_i2i_width = gr.Slider(label="Width", maximum=1024, value=512, step=8)
+                                    xl_i2i_height = gr.Slider(label="Height", maximum=1024, value=512, step=8)
+
+                                with gr.Column(scale=1):
+                                    xl_i2i_batch_size = gr.Slider(label="Batch Size", maximum=1, value=1)
+                                    xl_i2i_batch_count = gr.Slider(label="Batch Count", maximum=1, value=1)
+
+                            xl_i2i_cfg_scale = gr.Slider(label="CFG Scale", minimum=0.1, maximum=20, value=7, step=0.1)
+                            xl_i2i_denoising = gr.Slider(label="Denoising Strength", minimum=0, maximum=1, value=0.7, step=0.1)
+                            xl_i2i_seed = gr.Number(label="Seed", value=-1)
+                        
+
+                    with gr.Column(scale=2):
+                        xl_i2i_image_output = gr.Gallery(show_label=False, rows=2, preview=True)
+                        xl_send_to_png_i2i = gr.Button(value="Send INPUT IMAGE to PNG Info")
+                        xl_i2i_past_url = gr.Textbox(visible=False, interactive=False)
+
+                xl_i2i_text_button.click(xl_img2img, inputs=[xl_i2i_image_input, xl_i2i_denoising, xl_i2i_prompt, xl_i2i_negative_prompt,
+                                        xl_model, xl_i2i_steps, xl_i2i_sampler, xl_i2i_cfg_scale, xl_i2i_width, xl_i2i_height,
+                                        xl_i2i_seed], outputs=[xl_i2i_image_output, xl_i2i_past_url], concurrency_limit=64)
+        xl_send_to_img2img.click(send_to_img2img_def, inputs=xl_past_url, outputs=xl_i2i_image_input)
+        xl_send_to_png.click(send_to_img2img_def, inputs=xl_past_url, outputs=image_input)
+        xl_send_to_png_i2i.click(send_to_img2img_def, inputs=xl_i2i_past_url, outputs=image_input)
         with gr.Tab("HuggingFace Inference"):
             with gr.Row():
                 gr.Markdown("Add your model from HF.co, enter model ID.")