add ref image

app.py

@@ -11,19 +11,30 @@ device="cuda"
 model_id = "IDEA-CCNL/Taiyi-Stable-Diffusion-1B-Chinese-v0.1"
 
 pipe_text2img = StableDiffusionPipeline.from_pretrained(model_id).to(device)
+
+# pipe_inpaint = StableDiffusionInpaintPipeline.from_pretrained(model_id).to(device)    # work
+# pipe_inpaint = StableDiffusionInpaintPipeline(**pipe_text2img.components)   # not work
 pipe_img2img = StableDiffusionImg2ImgPipeline(**pipe_text2img.components)
-pipe_inpaint = StableDiffusionInpaintPipeline(**pipe_text2img.components)
 
-# io1_text2img = gr.Interface.load("models/IDEA-CCNL/Taiyi-Stable-Diffusion-1B-Chinese-v0.1")
+def infer_text2img(prompt, width, height): 
+    output = pipe_text2img(prompt, width=width, height=height, guidance_scale=7.5, num_inference_steps=20,)
+    image = output.images[0]
+    return image
+
+def infer_img2img(prompt, width, height, image_in, strength ):
+    init_image = image_in.convert("RGB").resize((width, height))
+    output = pipe_img2img(prompt, init_image=init_image, strength=strength, width=width, height=height, guidance_scale=7.5, num_inference_steps=20)
+    image = output.images[0]
+    return image
 
-def resize(w_val,l_val,img):
-  img = Image.open(img)
-  img = img.resize((w_val,l_val), Image.Resampling.LANCZOS)
-  return img
+def infer_inpaint(prompt, width, height, image_in): 
+    init_image = image_in["image"].convert("RGB").resize((width, height))
+    mask = image_in["mask"].convert("RGB").resize((width, height))
 
-def infer(prompt, guide, steps, width, height): 
-    output = pipe_text2img(prompt, guidance_scale=guide, num_inference_steps=steps, width=width, height=height)
-    # output = io1_text2img(prompt, guidance_scale=guide, num_inference_steps=steps, width=width, height=height)
+    output = pipe_inpaint(prompt, \
+                        init_image=init_image, mask_image=mask, \
+                        width=width, height=height, \
+                        guidance_scale=7.5, num_inference_steps=20)
     image = output.images[0]
     return image
 
@@ -35,20 +46,29 @@ with gr.Blocks() as demo:
                 ["女孩背影, 日落, 唯美插画"],
                 ]
     with gr.Row():
-        with gr.Column(scale=2, ):
-            output = gr.Image(label = '输出(output)')
-        
         with gr.Column(scale=1, ):
-            guide = gr.Slider(2, 15, value = 7, label = '文本引导强度(guidance scale)')
-            steps = gr.Slider(10, 30, value = 20, step = 1, label = '迭代次数(inference steps)')
+            # guide = gr.Slider(2, 15, value = 7, label = '文本引导强度(guidance scale)')
+            # steps = gr.Slider(10, 30, value = 20, step = 1, label = '迭代次数(inference steps)')
+            image_in = gr.Image(source='upload', elem_id="image_upload", type="pil", label="参考图")
             width = gr.Slider(256, 768, value = 512, step = 64, label = '宽度(width)')
             height = gr.Slider(256, 768, value = 512, step = 64, label = '高度(height)')
-            
+            strength = gr.Slider(0, 1.0, value = 0.8, step = 0.1, label = '参考图改变程度(strength)')
             prompt = gr.Textbox(label = '提示词(prompt)')
-            submit_btn = gr.Button("生成图片(Generate)").style(margin=False, rounded=(False, True, True, False), full_width=False,)
-
-            ex = gr.Examples(examples, fn=infer, inputs=[prompt, guide, steps, width, height], outputs=output)
-
-        submit_btn.click(fn = infer, inputs = [prompt, guide, steps, width, height], outputs = output)
+            submit_btn = gr.Button("生成图像(Generate)")
         
+        with gr.Column(scale=1, ):
+            image_out = gr.Image(label = '输出(output)')
+            ex = gr.Examples(examples, fn=infer_text2img, inputs=[prompt, width, height], outputs=image_out)
+        # with gr.Column(scale=1, ):
+        #     image_in = gr.Image(source='upload', tool='sketch', elem_id="image_upload", type="pil", label="Upload")
+        #     inpaint_prompt = gr.Textbox(label = '提示词(prompt)')
+        #     inpaint_btn = gr.Button("图像编辑(Inpaint)")
+            # img2img_prompt = gr.Textbox(label = '提示词(prompt)')
+            # img2img_btn = gr.Button("图像编辑(Inpaint)")
+        if isinstance(image_in, Image.Image):
+            submit_btn.click(fn = infer_img2img, inputs = [prompt, width, height, image_in, strength], outputs = image_out)
+        else:
+            submit_btn.click(fn = infer_text2img, inputs = [prompt, width, height], outputs = image_out)
+        # inpaint_btn.click(fn = infer_inpaint, inputs = [inpaint_prompt, width, height, image_in], outputs = image_out)
+        # img2img_btn.click(fn = infer_img2img, inputs = [img2img_prompt, width, height, image_in], outputs = image_out)
 demo.queue(concurrency_count=10).launch()