update app.py

app.py

@@ -4,7 +4,18 @@ import torch
 from PIL import Image
 import os
 
-
+scheduler = UniPCMultistepScheduler(
+    beta_start=0.00085,
+    beta_end=0.012,
+    beta_schedule="scaled_linear",
+    num_train_timesteps=1000,
+    trained_betas=None,
+    predict_epsilon=True,
+    thresholding=False,
+    predict_x0=True,
+    solver_type="bh2",
+    lower_order_final=True,
+)
 
 class Model:
     def __init__(self, name, path, prefix):
@@ -32,8 +43,8 @@ if torch.cuda.is_available():
   for model in models:
     try:
         unet = UNet2DConditionModel.from_pretrained(model.path, subfolder="unet", torch_dtype=torch.float16, use_auth_token=auth_token)
-        model.pipe_t2i = StableDiffusionPipeline.from_pretrained(model.path, unet=unet, vae=vae, torch_dtype=torch.float16, use_auth_token=auth_token)
-        model.pipe_i2i = StableDiffusionImg2ImgPipeline.from_pretrained(model.path, unet=unet, vae=vae, torch_dtype=torch.float16, use_auth_token=auth_token)
+        model.pipe_t2i = StableDiffusionPipeline.from_pretrained(model.path, unet=unet, vae=vae, torch_dtype=torch.float16, scheduler=scheduler, use_auth_token=auth_token)
+        model.pipe_i2i = StableDiffusionImg2ImgPipeline.from_pretrained(model.path, unet=unet, vae=vae, torch_dtype=torch.float16, scheduler=scheduler, use_auth_token=auth_token)
     except:
         models.remove(model)
   pipe = models[0].pipe_t2i
@@ -78,7 +89,6 @@ def txt_to_img(model_path, prompt, neg_prompt, guidance, steps, width, height, g
 
         pipe.to("cpu")
         pipe = current_model.pipe_t2i
-        pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
 
         if torch.cuda.is_available():
           pipe = pipe.to("cuda")
@@ -87,8 +97,12 @@ def txt_to_img(model_path, prompt, neg_prompt, guidance, steps, width, height, g
     prompt = current_model.prefix + prompt
     result = pipe(
       prompt,
+      negative_prompt = neg_prompt,
+      # num_images_per_prompt=n_images,
       num_inference_steps = int(steps),
       guidance_scale = guidance,
+      width = width,
+      height = height,
       generator = generator)
     
     return replace_nsfw_images(result)
@@ -103,7 +117,6 @@ def img_to_img(model_path, prompt, neg_prompt, img, strength, guidance, steps, w
 
         pipe.to("cpu")
         pipe = current_model.pipe_i2i
-        pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
         
         if torch.cuda.is_available():
           pipe = pipe.to("cuda")
@@ -224,7 +237,6 @@ with gr.Blocks(css=css) as demo:
           with gr.Tab("Options"):
             with gr.Group():
               neg_prompt = gr.Textbox(label="Negative prompt", placeholder="What to exclude from the image")
-              # neg_prompt = ""
 
               # n_images = gr.Slider(label="Images", value=1, minimum=1, maximum=4, step=1)
 
@@ -233,8 +245,8 @@ with gr.Blocks(css=css) as demo:
                 steps = gr.Slider(label="Steps", value=25, minimum=2, maximum=100, step=1)
 
               with gr.Row():
-                width = gr.Slider(label="Width", value=512, minimum=64, maximum=512, step=8)
-                height = gr.Slider(label="Height", value=512, minimum=64, maximum=512, step=8)
+                width = gr.Slider(label="Width", value=512, minimum=64, maximum=1024, step=8)
+                height = gr.Slider(label="Height", value=512, minimum=64, maximum=1024, step=8)
 
               seed = gr.Slider(0, 2147483647, label='Seed (0 = random)', value=0, step=1)
 
@@ -245,7 +257,7 @@ with gr.Blocks(css=css) as demo:
 
     # model_name.change(lambda x: gr.update(visible = x == models[0].name), inputs=model_name, outputs=custom_model_group)
 
-    inputs = [model_name, prompt, guidance, steps, width, height, seed, image, strength]
+    inputs = [model_name, prompt, guidance, steps, width, height, seed, image, strength, neg_prompt]
     prompt.submit(inference, inputs=inputs, outputs=image_out)
 
     generate.click(inference, inputs=inputs, outputs=image_out)
@@ -259,4 +271,4 @@ with gr.Blocks(css=css) as demo:
     ''')
 
 demo.queue(concurrency_count=1)
-demo.launch(debug=False, share=False)
+demo.launch(debug=False, share=False)