update

app.py

@@ -8,6 +8,32 @@ from modelscope.utils.constant import Tasks
 import cv2
 from diffusers import StableDiffusionXLPipeline
 import torch
+from diffusers import (
+    DiffusionPipeline,
+    DPMSolverMultistepScheduler,
+    DDIMScheduler,
+    HeunDiscreteScheduler,
+    EulerAncestralDiscreteScheduler,
+    EulerDiscreteScheduler,
+    PNDMScheduler
+)
+
+
+class KarrasDPM:
+    @staticmethod
+    def from_config(config):
+        return DPMSolverMultistepScheduler.from_config(config, use_karras_sigmas=True)
+
+
+SCHEDULERS = {
+    "DDIM": DDIMScheduler,
+    "DPMSolverMultistep": DPMSolverMultistepScheduler,
+    "HeunDiscrete": HeunDiscreteScheduler,
+    "KarrasDPM": KarrasDPM,
+    "K_EULER_ANCESTRAL": EulerAncestralDiscreteScheduler,
+    "K_EULER": EulerDiscreteScheduler,
+    "PNDM": PNDMScheduler,
+}
 
 
 deep_seek_llm = DeepSeekLLM(api_key=settings.deep_seek_api_key)
@@ -71,22 +97,26 @@ def bg_remove(_image, _type):
     return vis_img
 
 
-def text_to_image(_image, _prompt):
+def text_to_image(_prompt: str, _n_prompt: str, _scheduler: str, _inference_steps: int, _w: int, _h: int, _guidance_scale: float):
+    print('????????', _prompt, _scheduler, _inference_steps, _w, _h, _guidance_scale)
     t2i_pipeline = StableDiffusionXLPipeline.from_pretrained(
         "stabilityai/stable-diffusion-xl-base-1.0",
         torch_dtype=torch.float16,
         variant="fp16",
         use_safetensors=True,
     ).to("cuda")
-    result = t2i_pipeline(
-        prompt=_prompt,
-        negative_prompt='ugly',
-        num_inference_steps=22,
-        width=1024,
-        height=1024,
-        guidance_scale=7,
-    ).images[0]
-    return result
+    t2i_pipeline.scheduler = SCHEDULERS[_scheduler].from_config(t2i_pipeline.scheduler.config)
+    t2i_pipeline.enable_xformers_memory_efficient_attention()
+    with torch.inference_mode():
+        result = t2i_pipeline(
+            prompt=_prompt,
+            negative_prompt=_n_prompt,
+            num_inference_steps=_inference_steps,
+            width=_w,
+            height=_h,
+            guidance_scale=_guidance_scale,
+        ).images[0]
+        return result
 
 
 with gr.Blocks() as app:
@@ -254,10 +284,18 @@ with gr.Blocks() as app:
                 image = gr.Image()
 
             with gr.Column(scale=1, min_width=300):
-                with gr.Accordion(label="图像生成"):
+                with gr.Accordion(label="提示词", open=True):
                     prompt = gr.Textbox(label="提示语", value="", lines=3)
-                    t2i_btn = gr.Button('画图', variant='primary')
-            t2i_btn.click(fn=text_to_image, inputs=[prompt, image], outputs=[image])
+                    negative_prompt = gr.Textbox(label="负提示语", value="ugly", lines=2)
+                with gr.Accordion(label="参数设置", open=False):
+                    scheduler = gr.Dropdown(label='scheduler', choices=list(SCHEDULERS.keys()), value='KarrasDPM')
+                    inference_steps = gr.Number(label='inference steps', value=22, minimum=1, maximum=100)
+                    width = gr.Dropdown(label='width', choices=[512, 768, 832, 896, 1024, 1152], value=1024)
+                    height = gr.Dropdown(label='height', choices=[512, 768, 832, 896, 1024, 1152], value=1024)
+                    guidance_scale = gr.Number(label='guidance scale', value=7.0, minimum=1.0, maximum=10.0)
+                with gr.Row(variant='panel'):
+                    t2i_btn = gr.Button('🪄生成', variant='primary')
+            t2i_btn.click(fn=text_to_image, inputs=[prompt, negative_prompt, scheduler, inference_steps, width, height, guidance_scale], outputs=[image])
 
 
 app.launch(debug=settings.debug, show_api=False)

pyproject.toml

@@ -14,9 +14,9 @@ dependencies = [
     "langchain-openai>=0.1.16",
     "dashscope>=1.20.1",
     "modelscope[framework]>=1.16.1",
-    "torch==2.3.1",
-    "torchvision>=0.18.1",
-    "torchaudio>=2.3.1",
+    "torch==2.3.1+cu121",
+    "torchvision==0.18.1+cu121",
+    "torchaudio==2.3.1+cu121",
     "setuptools==69.5.1",
     "oss2>=2.18.6",
     "kornia>=0.7.3",
@@ -24,6 +24,7 @@ dependencies = [
     "tensorflow>=2.17.0",
     "transformers>=4.42.4",
     "accelerate>=0.32.1",
+    "xformers>=0.0.27",
 ]
 readme = "README.md"
 requires-python = ">= 3.8"
@@ -36,6 +37,11 @@ build-backend = "hatchling.build"
 managed = true
 dev-dependencies = []
 
+[[tool.rye.sources]]
+name = "torch"
+url = "https://download.pytorch.org/whl/cu121"
+type = "index"
+
 [tool.hatch.metadata]
 allow-direct-references = true
 

requirements-dev.lock

@@ -7,6 +7,7 @@
 #   all-features: false
 #   with-sources: false
 #   generate-hashes: false
+#   universal: false
 
 -e file:.
 absl-py==2.1.0
@@ -221,6 +222,7 @@ ml-dtypes==0.4.0
     # via tensorflow-intel
 modelscope==1.16.1
     # via aitoolkits-webui
+    # via modelscope
 mpmath==1.3.0
     # via sympy
 multidict==6.0.5
@@ -253,6 +255,7 @@ numpy==1.26.4
     # via tensorflow-intel
     # via torchvision
     # via transformers
+    # via xformers
 openai==1.35.13
     # via langchain-openai
 opencv-python==4.10.0.84
@@ -378,11 +381,6 @@ scipy==1.14.0
     # via modelscope
 semantic-version==2.10.0
     # via gradio
-setuptools==69.5.1
-    # via aitoolkits-webui
-    # via modelscope
-    # via tensorboard
-    # via tensorflow-intel
 shellingham==1.5.4
     # via typer
 simplejson==3.19.2
@@ -431,15 +429,16 @@ tomlkit==0.12.0
     # via gradio
 toolz==0.12.1
     # via altair
-torch==2.3.1
+torch==2.3.1+cu121
     # via accelerate
     # via aitoolkits-webui
     # via kornia
     # via torchaudio
     # via torchvision
-torchaudio==2.3.1
+    # via xformers
+torchaudio==2.3.1+cu121
     # via aitoolkits-webui
-torchvision==0.18.1
+torchvision==0.18.1+cu121
     # via aitoolkits-webui
 tqdm==4.66.4
     # via datasets
@@ -492,9 +491,16 @@ wheel==0.43.0
     # via astunparse
 wrapt==1.16.0
     # via tensorflow-intel
+xformers==0.0.27
+    # via aitoolkits-webui
 xxhash==3.4.1
     # via datasets
 yarl==1.9.4
     # via aiohttp
 zipp==3.19.2
     # via importlib-metadata
+setuptools==69.5.1
+    # via aitoolkits-webui
+    # via modelscope
+    # via tensorboard
+    # via tensorflow-intel

requirements.lock

@@ -7,6 +7,7 @@
 #   all-features: false
 #   with-sources: false
 #   generate-hashes: false
+#   universal: false
 
 -e file:.
 absl-py==2.1.0
@@ -221,6 +222,7 @@ ml-dtypes==0.4.0
     # via tensorflow-intel
 modelscope==1.16.1
     # via aitoolkits-webui
+    # via modelscope
 mpmath==1.3.0
     # via sympy
 multidict==6.0.5
@@ -253,6 +255,7 @@ numpy==1.26.4
     # via tensorflow-intel
     # via torchvision
     # via transformers
+    # via xformers
 openai==1.35.13
     # via langchain-openai
 opencv-python==4.10.0.84
@@ -378,11 +381,6 @@ scipy==1.14.0
     # via modelscope
 semantic-version==2.10.0
     # via gradio
-setuptools==69.5.1
-    # via aitoolkits-webui
-    # via modelscope
-    # via tensorboard
-    # via tensorflow-intel
 shellingham==1.5.4
     # via typer
 simplejson==3.19.2
@@ -431,15 +429,16 @@ tomlkit==0.12.0
     # via gradio
 toolz==0.12.1
     # via altair
-torch==2.3.1
+torch==2.3.1+cu121
     # via accelerate
     # via aitoolkits-webui
     # via kornia
     # via torchaudio
     # via torchvision
-torchaudio==2.3.1
+    # via xformers
+torchaudio==2.3.1+cu121
     # via aitoolkits-webui
-torchvision==0.18.1
+torchvision==0.18.1+cu121
     # via aitoolkits-webui
 tqdm==4.66.4
     # via datasets
@@ -492,9 +491,16 @@ wheel==0.43.0
     # via astunparse
 wrapt==1.16.0
     # via tensorflow-intel
+xformers==0.0.27
+    # via aitoolkits-webui
 xxhash==3.4.1
     # via datasets
 yarl==1.9.4
     # via aiohttp
 zipp==3.19.2
     # via importlib-metadata
+setuptools==69.5.1
+    # via aitoolkits-webui
+    # via modelscope
+    # via tensorboard
+    # via tensorflow-intel

requirements.txt

@@ -7,6 +7,7 @@
 #   all-features: false
 #   with-sources: false
 #   generate-hashes: false
+#   universal: false
 
 absl-py==2.1.0
     # via keras
@@ -220,6 +221,7 @@ ml-dtypes==0.4.0
     # via tensorflow-intel
 modelscope==1.16.1
     # via aitoolkits-webui
+    # via modelscope
 mpmath==1.3.0
     # via sympy
 multidict==6.0.5
@@ -252,6 +254,7 @@ numpy==1.26.4
     # via tensorflow-intel
     # via torchvision
     # via transformers
+    # via xformers
 openai==1.35.13
     # via langchain-openai
 opencv-python==4.10.0.84
@@ -377,11 +380,6 @@ scipy==1.14.0
     # via modelscope
 semantic-version==2.10.0
     # via gradio
-setuptools==69.5.1
-    # via aitoolkits-webui
-    # via modelscope
-    # via tensorboard
-    # via tensorflow-intel
 shellingham==1.5.4
     # via typer
 simplejson==3.19.2
@@ -416,7 +414,6 @@ tensorboard-data-server==0.7.2
     # via tensorboard
 tensorflow==2.17.0
     # via aitoolkits-webui
-    # via tensorflow
 tensorflow-io-gcs-filesystem==0.31.0
     # via tensorflow-intel
 termcolor==2.4.0
@@ -429,15 +426,16 @@ tomlkit==0.12.0
     # via gradio
 toolz==0.12.1
     # via altair
-torch==2.3.1
+torch==2.3.1+cu121
     # via accelerate
     # via aitoolkits-webui
     # via kornia
     # via torchaudio
     # via torchvision
-torchaudio==2.3.1
+    # via xformers
+torchaudio==2.3.1+cu121
     # via aitoolkits-webui
-torchvision==0.18.1
+torchvision==0.18.1+cu121
     # via aitoolkits-webui
 tqdm==4.66.4
     # via datasets
@@ -490,9 +488,16 @@ wheel==0.43.0
     # via astunparse
 wrapt==1.16.0
     # via tensorflow-intel
+xformers==0.0.27
+    # via aitoolkits-webui
 xxhash==3.4.1
     # via datasets
 yarl==1.9.4
     # via aiohttp
 zipp==3.19.2
     # via importlib-metadata
+setuptools==69.5.1
+    # via aitoolkits-webui
+    # via modelscope
+    # via tensorboard
+    # via tensorflow-intel