Skip to content

Commit 75ad392

Browse files
saddam213saddam213
committed
QwenImageEdit support
1 parent 8945371 commit 75ad392

File tree

10 files changed

+276
-80
lines changed

10 files changed

+276
-80
lines changed

‎TensorStack.Python/Common/PipelineOptions.cs‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
using System.Collections.Generic;
2+
using TensorStack.Common.Tensor;
23

34
namespace TensorStack.Python.Common
45
{
@@ -21,6 +22,6 @@ public record PipelineOptions
2122
public SchedulerType Scheduler { get; set; }
2223
public SchedulerType[] Schedulers { get; set; }
2324
public List<LoraOptions> LoraOptions { get; set; }
24-
25+
public ImageTensor ImageInput { get; set; }
2526
}
2627
}

‎TensorStack.Python/Common/ProcessType.cs‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
public enum ProcessType
44
{
55
TextToImage = 0,
6-
ImageToImage = 1
6+
ImageToImage = 1,
7+
ImageEdit = 2
78
}
89
}

‎TensorStack.Python/Config/EnvironmentConfig.cs‎

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,5 +6,86 @@ public record EnvironmentConfig
66
public string Environment { get; set; }
77
public string[] Requirements { get; set; }
88
public string Directory { get; set; }
9+
10+
11+
public readonly static EnvironmentConfig DefaultCPU = new()
12+
{
13+
Environment = "default-cpu",
14+
Directory = "PythonRuntime",
15+
Requirements =
16+
[
17+
"torchvision==0.22.0",
18+
19+
// Default Packages
20+
"typing",
21+
"wheel",
22+
"transformers",
23+
"accelerate",
24+
"diffusers",
25+
"protobuf",
26+
"sentencepiece",
27+
"pillow",
28+
"ftfy",
29+
"scipy",
30+
"peft",
31+
"pillow"
32+
]
33+
};
34+
35+
36+
public readonly static EnvironmentConfig DefaultCUDA = new()
37+
{
38+
Environment = "default-cuda",
39+
Directory = "PythonRuntime",
40+
Requirements =
41+
[
42+
"--extra-index-url https://download.pytorch.org/whl/cu118",
43+
"torchvision==0.22.0+cu118",
44+
45+
// Default Packages
46+
"typing",
47+
"wheel",
48+
"transformers",
49+
"accelerate",
50+
"diffusers",
51+
"protobuf",
52+
"sentencepiece",
53+
"pillow",
54+
"ftfy",
55+
"scipy",
56+
"peft",
57+
"pillow"
58+
]
59+
};
60+
61+
62+
public readonly static EnvironmentConfig DefaultROCM = new()
63+
{
64+
Environment = "default-rocm",
65+
Directory = "PythonRuntime",
66+
Requirements =
67+
[
68+
"https://repo.radeon.com/rocm/windows/rocm-rel-7.1.1/rocm_sdk_core-0.1.dev0-py3-none-win_amd64.whl",
69+
"https://repo.radeon.com/rocm/windows/rocm-rel-7.1.1/rocm_sdk_devel-0.1.dev0-py3-none-win_amd64.whl",
70+
"https://repo.radeon.com/rocm/windows/rocm-rel-7.1.1/rocm_sdk_libraries_custom-0.1.dev0-py3-none-win_amd64.whl",
71+
"https://repo.radeon.com/rocm/windows/rocm-rel-7.1.1/rocm-0.1.dev0.tar.gz",
72+
"https://repo.radeon.com/rocm/windows/rocm-rel-7.1.1/torch-2.9.0+rocmsdk20251116-cp312-cp312-win_amd64.whl",
73+
"https://repo.radeon.com/rocm/windows/rocm-rel-7.1.1/torchvision-0.24.0+rocmsdk20251116-cp312-cp312-win_amd64.whl",
74+
75+
// Default Packages
76+
"typing",
77+
"wheel",
78+
"transformers",
79+
"accelerate",
80+
"diffusers",
81+
"protobuf",
82+
"sentencepiece",
83+
"pillow",
84+
"ftfy",
85+
"scipy",
86+
"peft",
87+
"pillow"
88+
]
89+
};
990
}
1091
}

‎TensorStack.Python/Pipelines/ChromaPipeline.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
from collections.abc import Buffer
88
from typing import Coroutine, Dict, Sequence, List, Tuple, Optional
99
from diffusers import ChromaPipeline, ChromaImg2ImgPipeline
10-
from tensorstack.utils import MemoryStdout, create_scheduler, getDataType, createTensor
10+
from tensorstack.utils import MemoryStdout, create_scheduler, getDataType, imageFromInput
1111
sys.stderr = MemoryStdout()
1212

1313
# Globals
@@ -150,7 +150,7 @@ def generate(
150150
)[0]
151151
elif _processType == "ImageToImage":
152152
output = _pipeline(
153-
image = createTensor(inputData, inputShape, device=_pipeline.device, dtype=_pipeline.dtype),
153+
image = imageFromInput(inputData, inputShape),
154154
strength = strength,
155155
prompt = prompt,
156156
negative_prompt = negativePrompt,

‎TensorStack.Python/Pipelines/QwenImagePipeline.py‎

Lines changed: 26 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,8 @@
66
from threading import Event
77
from collections.abc import Buffer
88
from typing import Coroutine, Dict, Sequence, List, Tuple, Optional
9-
from diffusers import QwenImagePipeline, QwenImageImg2ImgPipeline
10-
from tensorstack.utils import MemoryStdout, create_scheduler, getDataType, createTensor
9+
from diffusers import QwenImagePipeline, QwenImageImg2ImgPipeline, QwenImageEditPipeline
10+
from tensorstack.utils import MemoryStdout, create_scheduler, getDataType, imageFromInput
1111
sys.stderr = MemoryStdout()
1212

1313
# Globals
@@ -56,6 +56,14 @@ def load(
5656
token = secureToken,
5757
variant=variant
5858
)
59+
elif _processType == "ImageEdit":
60+
_pipeline = QwenImageEditPipeline.from_pretrained(
61+
modelName,
62+
torch_dtype=torch_dtype,
63+
cache_dir = cacheDir,
64+
token = secureToken,
65+
variant=variant
66+
)
5967

6068
#Lora Adapters
6169
if loraAdapters is not None:
@@ -151,7 +159,7 @@ def generate(
151159
)[0]
152160
elif _processType == "ImageToImage":
153161
output = _pipeline(
154-
image = createTensor(inputData, inputShape, device=_pipeline.device, dtype=_pipeline.dtype),
162+
image = imageFromInput(inputData, inputShape),
155163
strength = strength,
156164
prompt = prompt,
157165
negative_prompt = negativePrompt,
@@ -165,6 +173,21 @@ def generate(
165173
callback_on_step_end = _progress_callback,
166174
callback_on_step_end_tensor_inputs = ["latents"]
167175
)[0]
176+
elif _processType == "ImageEdit":
177+
output = _pipeline(
178+
image = imageFromInput(inputData, inputShape),
179+
prompt = prompt,
180+
negative_prompt = negativePrompt,
181+
height = height,
182+
width = width,
183+
generator = _generator.manual_seed(seed),
184+
true_cfg_scale = guidanceScale,
185+
guidance_scale = guidanceScale2,
186+
num_inference_steps = steps,
187+
output_type = "np",
188+
callback_on_step_end = _progress_callback,
189+
callback_on_step_end_tensor_inputs = ["latents"]
190+
)[0]
168191

169192
# (Batch, Channel, Height, Width)
170193
output = output.transpose(0, 3, 1, 2)

‎TensorStack.Python/Pipelines/WanPipeline.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
from collections.abc import Buffer
88
from typing import Coroutine, Dict, Sequence, List, Tuple, Optional
99
from diffusers import WanPipeline, WanImageToVideoPipeline, UniPCMultistepScheduler
10-
from tensorstack.utils import MemoryStdout, create_scheduler, getDataType, createTensor
10+
from tensorstack.utils import MemoryStdout, create_scheduler, getDataType, tensorFromInput
1111
sys.stderr = MemoryStdout()
1212

1313
# Globals
@@ -152,7 +152,7 @@ def generate(
152152
)[0]
153153
elif _processType == "ImageToImage":
154154
max_area = height * width
155-
image = createTensor(inputData, inputShape, device=_pipeline.device, dtype=_pipeline.dtype)
155+
image = tensorFromInput(inputData, inputShape, device=_pipeline.device, dtype=_pipeline.dtype)
156156
image = resize_tensor(image, max_area, _pipeline)
157157
output = _pipeline(
158158
image = image,

‎TensorStack.Python/Pipelines/ZImagePipeline.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
from collections.abc import Buffer
88
from typing import Coroutine, Dict, Sequence, List, Tuple, Optional
99
from diffusers import ZImagePipeline, ZImageImg2ImgPipeline
10-
from tensorstack.utils import MemoryStdout, create_scheduler, getDataType, createTensor
10+
from tensorstack.utils import MemoryStdout, create_scheduler, getDataType, imageFromInput
1111
sys.stderr = MemoryStdout()
1212

1313
# Globals
@@ -150,7 +150,7 @@ def generate(
150150
)[0]
151151
elif _processType == "ImageToImage":
152152
output = _pipeline(
153-
image = createTensor(inputData, inputShape, device=_pipeline.device, dtype=_pipeline.dtype),
153+
image = imageFromInput(inputData, inputShape),
154154
strength = strength,
155155
prompt = prompt,
156156
negative_prompt = negativePrompt,

‎TensorStack.Python/Python/lib/tensorstack/utils.py‎

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import torch
2+
from PIL import Image
23
from typing import Sequence, Optional
34
import numpy as np
45
import threading
@@ -80,7 +81,7 @@ def getDataType(dtype: str):
8081
return torch.float
8182

8283

83-
def createTensor(
84+
def tensorFromInput(
8485
inputData: Optional[Sequence[float]],
8586
inputShape: Optional[Sequence[int]],
8687
*,
@@ -112,6 +113,23 @@ def createTensor(
112113
return torch.from_numpy(np_array).to(device=device, dtype=dtype)
113114

114115

116+
def imageFromInput(
117+
inputData: Optional[Sequence[float]],
118+
inputShape: Optional[Sequence[int]],
119+
) -> Optional[Image.Image]:
120+
121+
if not inputData or not inputShape:
122+
return None
123+
124+
t = torch.tensor(inputData, dtype=torch.float32)
125+
t = t.view(*inputShape)
126+
t = t[0]
127+
t = (t + 1) / 2
128+
t = t.permute(1, 2, 0)
129+
t = (t.clamp(0, 1) * 255).to(torch.uint8)
130+
return Image.fromarray(t.numpy())
131+
132+
115133
class MemoryStdout:
116134
def __init__(self, callback=None):
117135
self.callback = callback

0 commit comments

Comments
 (0)