Update app.py

app.py

@@ -49,6 +49,7 @@ def tpu_inference_api(
     do_img2img: bool = False,
     init_image: Optional[str] = None,
     image2image_strength: float = 0,
+    inference_steps = 25,
 ) -> bytes:
     url = os.getenv("TPU_INFERENCE_API")
     if(randomize_seed):
@@ -67,7 +68,8 @@ def tpu_inference_api(
         "seed": seed,
         "do_img2img": do_img2img,
         "image2image_strength": image2image_strength,
-        "init_image": init_image
+        "init_image": init_image,
+        "inference_steps": inference_steps,
     }
     response = requests.post(url, json=payload)
     if response.status_code != 200:
@@ -155,7 +157,7 @@ def run(prompt, radio="model-v2", preset=PRESET_Q, h=1216, w=832, negative_promp
     return zero_inference_api(prompt, radio, preset, h, w, negative_prompt, guidance_scale, randomize_seed, seed, do_img2img, init_image, image2image_strength)
     
 @spaces.GPU
-def zero_inference_api(prompt, radio="model-v2", preset=PRESET_Q, h=1216, w=832, negative_prompt=NEGATIVE_PROMPT, guidance_scale=4.0, randomize_seed=True, seed=42, do_img2img=False, init_image=None, image2image_strength=0, progress=gr.Progress(track_tqdm=True)):
+def zero_inference_api(prompt, radio="model-v2", preset=PRESET_Q, h=1216, w=832, negative_prompt=NEGATIVE_PROMPT, guidance_scale=4.0, randomize_seed=True, seed=42, do_img2img=False, init_image=None, image2image_strength=0, inference_steps=25, progress=gr.Progress(track_tqdm=True)):
     prompt = prompt.strip() + ", " + preset.strip()
     negative_prompt = negative_prompt.strip() if negative_prompt and negative_prompt.strip() else None
     
@@ -167,24 +169,26 @@ def zero_inference_api(prompt, radio="model-v2", preset=PRESET_Q, h=1216, w=832,
         guidance_scale = 0.0
         
     generator = torch.Generator(device="cuda").manual_seed(seed)
+    if inference_steps > 50:
+        inference_steps = 50
     
     if not do_img2img:
         if radio == "model-v2":
-            image = pipe(prompt, height=h, width=w, negative_prompt=negative_prompt, guidance_scale=guidance_scale, guidance_rescale=0.75, generator=generator, num_inference_steps=25).images[0]    
+            image = pipe(prompt, height=h, width=w, negative_prompt=negative_prompt, guidance_scale=guidance_scale, guidance_rescale=0.75, generator=generator, num_inference_steps=inference_steps).images[0]    
         else:
-            image = pipe2(prompt, height=h, width=w, negative_prompt=negative_prompt, guidance_scale=guidance_scale, guidance_rescale=0.75, generator=generator, num_inference_steps=25).images[0]           
+            image = pipe2(prompt, height=h, width=w, negative_prompt=negative_prompt, guidance_scale=guidance_scale, guidance_rescale=0.75, generator=generator, num_inference_steps=inference_steps).images[0]           
     else:
         init_image = Image.fromarray(init_image)
         if radio == "model-v2":
-            image = pipe_img2img(prompt, image=init_image, strength=image2image_strength, negative_prompt=negative_prompt, guidance_scale=guidance_scale, generator=generator, num_inference_steps=20).images[0]    
+            image = pipe_img2img(prompt, image=init_image, strength=image2image_strength, negative_prompt=negative_prompt, guidance_scale=guidance_scale, generator=generator, num_inference_steps=inference_steps).images[0]    
         else:
-            image = pipe2_img2img(prompt, image=init_image, strength=image2image_strength, negative_prompt=negative_prompt, guidance_scale=guidance_scale, generator=generator, num_inference_steps=20).images[0]
+            image = pipe2_img2img(prompt, image=init_image, strength=image2image_strength, negative_prompt=negative_prompt, guidance_scale=guidance_scale, generator=generator, num_inference_steps=inference_steps).images[0]
     
     naifix = prompt[:40].replace(":", "_").replace("\\", "_").replace("/", "_") + f" s-{seed}-"
     with tempfile.NamedTemporaryFile(prefix=naifix, suffix=".png", delete=False) as tmpfile:
         parameters = {
             "prompt": prompt,
-            "steps": 25, 
+            "steps": inference_steps, 
             "height": h,
             "width": w,
             "scale": guidance_scale,
@@ -247,6 +251,7 @@ with gr.Blocks(theme=theme) as demo:
 
                 preset = gr.Textbox(show_label=False, scale=5, value=PRESET_Q, info="Quality presets")
                 radio = gr.Radio(["model-v2-beta", "model-v2"], value="model-v2", label = "Choose the inference model")
+                inference_steps = gr.Slider(label="Inference Steps", value=25, minimum=4, maximum=50, step=1)
                 with gr.Row():                    
                     height = gr.Slider(label="Height", value=1216, minimum=512, maximum=2560, step=64)
                     width = gr.Slider(label="Width", value=832, minimum=512, maximum=2560, step=64)
@@ -262,7 +267,7 @@ with gr.Blocks(theme=theme) as demo:
             image2image_resize = gr.Checkbox(label="Resize input image", value=False, visible=False)
             image2image_strength = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, label="Noising strength", value=0.7, visible=False)
                 
-        with gr.Column():
+        with gr.Column(scale=2.5):
             output = gr.Image(type="filepath", interactive=False)
 
     gr.Examples(fn=run, examples=["mayano_top_gun_\(umamusume\), 1girl, rurudo", "sho (sho lwlw),[[[ohisashiburi]]],fukuro daizi,tianliang duohe fangdongye,[daidai ookami],year_2023, (wariza), depth of field, official_art"], inputs=prompt, outputs=[output, seed], cache_examples="lazy")
@@ -278,8 +283,9 @@ with gr.Blocks(theme=theme) as demo:
             prompt.submit
         ],
         fn=run,
-        inputs=[prompt, radio, preset, height, width, negative_prompt, guidance_scale, randomize_seed, seed, tpu_inference, do_img2img, init_image, image2image_resize, image2image_strength],
+        inputs=[prompt, radio, preset, height, width, negative_prompt, guidance_scale, randomize_seed, seed, tpu_inference, do_img2img, init_image, image2image_resize, image2image_strength, inference_steps],
         outputs=[output, seed],
+        concurrency_limit=1,
     )
 if __name__ == "__main__":
     demo.launch(share=True)