add txt2img

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
 
 def empty_checker(images, **kwargs): return images, False
@@ -19,10 +19,16 @@ YOUR_TOKEN=MY_SECRET_TOKEN
 
 device="cpu"
 
+# img2img pipeline
 img_pipe = StableDiffusionImg2ImgPipeline.from_pretrained("AkiKagura/mkgen-diffusion", duse_auth_token=YOUR_TOKEN)
 img_pipe.safety_checker = empty_checker
 img_pipe.to(device)
 
+# txt2img pipeline
+pipe = StableDiffusionPipeline.from_pretrained("AkiKagura/mkgen-diffusion", duse_auth_token=YOUR_TOKEN)
+pipe.safety_checker = empty_checker
+pipe.to(device)
+
 source_img = gr.Image(source="upload", type="filepath", label="init_img | 512*512 px")
 gallery = gr.Gallery(label="Generated images", show_label=False, elem_id="gallery").style(grid=[1], height="auto")
 
@@ -42,7 +48,10 @@ def infer(source_img, prompt, guide, steps, seed, strength):
     source_image = resize(512, source_img)
     source_image.save('source.png')
     
-    images_list = img_pipe([prompt] * 1, init_image=source_image, strength=strength, guidance_scale=guide, num_inference_steps=steps)
+    if source_image is None:
+        images_list = pipe([prompt] * 1, guidance_scale=guide, num_inference_steps=steps)
+    else:
+        images_list = img_pipe([prompt] * 1, init_image=source_image, strength=strength, guidance_scale=guide, num_inference_steps=steps)
     images = []
     
     for i, image in enumerate(images_list["images"]):