Update template/app.py

template/app.py

@@ -19,29 +19,19 @@ scheduler = DPMSolverMultistepScheduler(
     lower_order_final=True,
 )
 
-last_mode = "txt2img"
+pipe = StableDiffusionPipeline.from_pretrained(
+  model_name,
+  torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+  scheduler=scheduler)
 
-def get_pipe(img2img=False):
+pipe_i2i = StableDiffusionImg2ImgPipeline.from_pretrained(
+  model_name,
+  torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+  scheduler=scheduler)
 
-    if img2img:
-      pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
-        model_name,
-        torch_dtype=torch.float16 if torch.cuda.is_available() else 'auto',
-        scheduler=scheduler)
-    else:
-      pipe = StableDiffusionPipeline.from_pretrained(
-        model_name,
-        torch_dtype=torch.float16 if torch.cuda.is_available() else 'auto',
-        scheduler=scheduler)
-    
-    if torch.cuda.is_available():
-      pipe = pipe.to("cuda")
-    
-    return pipe
-
-pipe = get_pipe()
-
-device = "GPU 🔥" if torch.cuda.is_available() else "CPU 🥶"
+if torch.cuda.is_available():
+  pipe = pipe.to("cuda")
+  pipe_i2i = pipe_i2i.to("cuda")
 
 def error_str(error, title="Error"):
     return f"""#### {title}
@@ -62,14 +52,6 @@ def inference(prompt, guidance, steps, width=512, height=512, seed=0, img=None,
 
 def txt_to_img(prompt, neg_prompt, guidance, steps, width, height, generator):
 
-    global last_mode
-    global pipe
-    if last_mode != "txt2img":
-
-        pipe = get_pipe()
-        last_mode = "txt2img"
-
-    prompt = prompt  
     result = pipe(
       prompt,
       negative_prompt = neg_prompt,
@@ -83,17 +65,9 @@ def txt_to_img(prompt, neg_prompt, guidance, steps, width, height, generator):
 
 def img_to_img(prompt, neg_prompt, img, strength, guidance, steps, width, height, generator):
 
-    global last_mode
-    global pipe
-    if last_mode != "img2img":
-       
-        pipe = get_pipe(img2img=True)
-        last_mode = "img2img"
-
-    prompt = prompt
     ratio = min(height / img.height, width / img.width)
     img = img.resize((int(img.width * ratio), int(img.height * ratio)), Image.LANCZOS)
-    result = pipe(
+    result = pipe_i2i(
         prompt,
         negative_prompt = neg_prompt,
         init_image = img,
@@ -126,6 +100,7 @@ with gr.Blocks(css=css) as demo:
                Demo for the <a href="https://huggingface.co/$model_name">$model_name</a> Stable Diffusion model.<br>
                Add the following tokens to your prompts for the effect: <b>$prefix</b>.
               </p>
+              Running on <b>{"GPU 🔥" if torch.cuda.is_available() else "CPU 🥶"}</b>
             </div>
         """
     )