Update app.py

app.py

@@ -12,6 +12,11 @@ pipe = StableDiffusionPipeline.from_pretrained(
   model_id,
   torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
   scheduler=scheduler)
+    
+    
+    pipe = StableDiffusionPipeline.from_pretrained("hakurei/waifu-diffusion", use_auth_token="<mytoken>").to('cuda')
+    image = pipe(prompt, safety_checker=None).images[0]
+    image.save(f"{prompt}.png")
 
 pipe_i2i = StableDiffusionImg2ImgPipeline.from_pretrained(
   model_id,
@@ -74,7 +79,7 @@ def txt_to_img(prompt, neg_prompt, guidance, steps, width, height, generator):
       height = height,
       generator = generator)
     
-    return result.images[0]
+    return result.imagesimage = pipe(prompt, safety_checker=None).images[0]
 
 def img_to_img(prompt, neg_prompt, img, strength, guidance, steps, width, height, generator):
 
@@ -131,10 +136,10 @@ Have Fun & Enjoy
         with gr.Column(scale=55):
           with gr.Group():
               with gr.Row():
-                prompt = gr.Textbox(label="Prompt", show_label=False, max_lines=2,placeholder=f"{prefix} [your prompt]").style(container=False)
+                prompt = gr.Textbox(label="Prompt", show_label=False,max_lines=2,placeholder=f"{prefix} [your prompt]").style(container=False)
                 generate = gr.Button(value="Generate").style(rounded=(False, True, True, False))
 
-              image_out = gr.Image(height=512)
+              image_out = gr.Image(image = pipe (prompt, safety_checker=None) height=512)(f"{prompt}.png")
           error_output = gr.Markdown()
 
         with gr.Column(scale=45):