Update app.py

app.py

@@ -8,7 +8,7 @@ from io import BytesIO
 import os
 MY_SECRET_TOKEN=os.environ.get('HF_TOKEN_SD')
 
-#from diffusers import StableDiffusionPipeline
+from diffusers import StableDiffusionPipeline
 from diffusers import StableDiffusionImg2ImgPipeline
 
 print("hello sylvain")
@@ -17,9 +17,6 @@ YOUR_TOKEN=MY_SECRET_TOKEN
 
 device="cpu"
 
-pipe = StableDiffusionImg2ImgPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", use_auth_token=YOUR_TOKEN)
-pipe.to(device)
-
 source_img = gr.Image(source="upload", type="filepath", label="init_img")
 gallery = gr.Gallery(label="Generated images", show_label=False, elem_id="gallery").style(grid=[2], height="auto")
 
@@ -33,9 +30,17 @@ def resize(width,img):
 
 
 def infer(prompt, source_img): 
-    init_image = resize(512,source_img)
-    init_image.save('source.png')
-    images_list = pipe([prompt] * 2, init_image=init_image, strength=0.75)
+    if(source_img != None):
+        pipe = StableDiffusionImg2ImgPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", use_auth_token=YOUR_TOKEN)
+        pipe.to(device)
+        init_image = resize(512,source_img)
+        init_image.save('source.png')
+        images_list = pipe([prompt] * 2, init_image=init_image, strength=0.75)
+    else:
+        pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", use_auth_token=YOUR_TOKEN)
+        pipe.to(device)
+        images_list = pipe([prompt] * 2)
+    
     images = []
     safe_image = Image.open(r"unsafe.png")
     for i, image in enumerate(images_list["sample"]):