Upload main.py

main.py

@@ -10,7 +10,16 @@ import retrying
 from PIL import Image
 
 @retrying.retry(stop_max_attempt_number=2, wait_fixed=1000)
-def generate_image(image, inference_steps, guidance_scale, image_size, prompt, model_choice):
+def generate_image(image, inference_steps, guidance_scale, image_size, prompt, model_choice, style_name):
+    # 检查模型名称是否包含SDXL-Lightning
+    if "SDXL-Lightning" in model_choice:
+        inference_steps = 4
+        guidance_scale = 1
+
+    if "turbo" in model_choice:
+        inference_steps = 6
+        guidance_scale = 1
+
     image_size = {
         '1:1': '1024x1024',
         '1:2': '1024x2048',
@@ -44,6 +53,10 @@ def generate_image(image, inference_steps, guidance_scale, image_size, prompt, m
         # 如果选择了图生图模型但没有提供图像，可以返回一个错误消息或使用默认图像
         return Image.new('RGB', (512, 512), color = 'white')  # 创建一个白色的默认图像
 
+    if "PhotoMaker" in model_choice:
+        data["style_name"] = style_name
+        data["style_strengh_radio"] = 20
+
     url = f"https://api.siliconflow.cn/v1/{model_choice}"
     
     with requests.post(url, data=json.dumps(data), headers=headers) as response:
@@ -53,40 +66,52 @@ def generate_image(image, inference_steps, guidance_scale, image_size, prompt, m
             img = Image.open(io.BytesIO(image_response.content))
             return img
 
-def dosomething(image, inference_steps, guidance_scale, image_size, prompt, model_choice):
-    return generate_image(image, inference_steps, guidance_scale, image_size, prompt, model_choice)
+def dosomething(image, inference_steps, guidance_scale, image_size, prompt, model_choice, style_name):
+    return generate_image(image, inference_steps, guidance_scale, image_size, prompt, model_choice, style_name)
 
 if __name__ == "__main__":
     dotenv.load_dotenv()
     
     with gr.Blocks() as demo:
         gr.Markdown("# 文生图、图生图")
-        
+
+        guidance_scale = gr.Slider(minimum=0, maximum=100, value=7.5, step=0.1, label='引导比例')
         with gr.Row():
             inference_steps = gr.Number(minimum=1, maximum=100, value=20, label='推理步数')
-            guidance_scale = gr.Slider(minimum=0, maximum=100, value=7.5, step=0.1, label='引导比例')
             image_size = gr.Dropdown(['1:1', '1:2', '3:2', '3:4', '16:9', '9:16'], label='生成图片比例', value='1:1')
+            style_name = gr.Dropdown([
+                                       'Cinematic',
+                                       'Comic book',
+                                       'Disney Character',
+                                       'Digital Art',
+                                       'Photographic (Default)',
+                                       'Fantasy Art',
+                                       'Neopunk',
+                                       'Enhance',
+                                       'Lowpoly',
+                                       'Line art',
+                                       '(No style)'
+                                       ], label='PhotoMaker风格', value='Photographic (Default)')
+
+            model_choice = gr.Dropdown([
+                                       'stabilityai/stable-diffusion-3-medium/text-to-image',
+                                       'stabilityai/stable-diffusion-xl-base-1.0/text-to-image',
+                                       'stabilityai/stable-diffusion-2-1/text-to-image',
+                                       'stabilityai/sd-turbo/text-to-image',
+                                       'stabilityai/sdxl-turbo/text-to-image',
+                                       'ByteDance/SDXL-Lightning/text-to-image',
+                                       'stabilityai/stable-diffusion-xl-base-1.0/image-to-image',
+                                       'stabilityai/stable-diffusion-2-1/image-to-image',
+                                       'ByteDance/SDXL-Lightning/image-to-image',
+                                       'TencentARC/PhotoMaker/image-to-image',
+                                       ],label='选择模型', value='stabilityai/stable-diffusion-xl-base-1.0/image-to-image')
+
+        prompt_input = gr.Textbox(label="输入prompt", lines=3,  value="Transform all objects in the scene into a highly detailed and realistic anime style. Ensure that all characters have perfectly proportioned features including complete and natural-looking hands and fingers, and symmetrical, well-defined facial features with no distortions or anomalies. All objects should be rendered with vibrant and colorful details, smooth shading, and dynamic compositions. The style should resemble the works of Studio Ghibli or Makoto Shinkai, with meticulous attention to detail in every aspect, including backgrounds, clothing, and accessories. The overall image should be cohesive, with a harmonious blend of all elements.", placeholder="在此输入你的prompt...")
 
         with gr.Row():
-            prompt_input = gr.Textbox(label="输入prompt", lines=3, placeholder="在此输入你的prompt...")
-            with gr.Column():
-                model_choice = gr.Dropdown([
-                                           'stabilityai/stable-diffusion-3-medium/text-to-image',
-                                           'stabilityai/stable-diffusion-xl-base-1.0/text-to-image',
-                                           'stabilityai/stable-diffusion-2-1/text-to-image', 
-                                           'stabilityai/sd-turbo/text-to-image', 
-                                           'stabilityai/sdxl-turbo/text-to-image', 
-                                           'ByteDance/SDXL-Lightning/text-to-image', 
-                                           'stabilityai/stable-diffusion-xl-base-1.0/image-to-image', 
-                                           'stabilityai/stable-diffusion-2-1/image-to-image', 
-                                           'stabilityai/SDXL-Lightning/image-to-image', 
-                                           'InstantX/InstantID/image-to-image', 
-                                           'TencentARC/PhotoMaker/image-to-image'
-], 
-                                           label='选择模型', value='stabilityai/stable-diffusion-xl-base-1.0/image-to-image')
-                with gr.Row():
-                    clear_btn = gr.Button("清空")
-                    generate_btn = gr.Button("生成图片")
+            #with gr.Column():
+            clear_btn = gr.Button("清空")
+            generate_btn = gr.Button("生成图片")
 
         with gr.Row():
             image_input = gr.Image(label='选择原图 (仅用于图生图)', type='pil')
@@ -94,14 +119,14 @@ if __name__ == "__main__":
 
         generate_btn.click(
             fn=dosomething,
-            inputs=[image_input, inference_steps, guidance_scale, image_size, prompt_input, model_choice],
+            inputs=[image_input, inference_steps, guidance_scale, image_size, prompt_input, model_choice, style_name],
             outputs=output_image
         )
         
         clear_btn.click(
-            fn=lambda: (None, 20, 7.5, "1:1", "", "stable-diffusion-xl-base-1.0/image-to-image", None),
+            fn=lambda: (None, 20, 5, "1:1", "a half-body portrait of a man img wearing the sunglasses in Iron man suit, best quality", "TencentARC/PhotoMaker/image-to-image", "Photographic (Default)", None),
             inputs=None,
-            outputs=[image_input, inference_steps, guidance_scale, image_size, prompt_input, model_choice, output_image]
+            outputs=[image_input, inference_steps, guidance_scale, image_size, prompt_input, model_choice, style_name, output_image]
         )
 
     demo.launch(server_name='0.0.0.0', server_port=7861)