Skip to content

Commit 7600872

Browse files
MeatfuckerMeatfucker
committed
Turned a bunch of repetitive boilerplate into functions
1 parent 95315cc commit 7600872

43 files changed

Lines changed: 456 additions & 1238 deletions

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎avernus.py‎

Lines changed: 129 additions & 507 deletions
Large diffs are not rendered by default.

‎modules/auraflow.py‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
import torch
66

77
from pydantic_models import AuraFlowRequest, ImageResponse
8-
from utils import image_to_base64
8+
from utils import image_to_base64, get_seed_generators, load_loras
99

1010
PIPELINE: AuraFlowPipeline
1111
LOADED: bool = False
@@ -15,15 +15,10 @@
1515

1616
def load_auraflow_pipeline(model_name="fal/AuraFlow"):
1717
global PIPELINE
18-
PIPELINE = AuraFlowPipeline.from_pretrained(model_name, torch_dtype=dtype).to("cuda")
18+
PIPELINE = AuraFlowPipeline.from_pretrained(model_name, torch_dtype=dtype, safety_checker=None).to("cuda")
1919
PIPELINE.enable_model_cpu_offload()
2020
PIPELINE.vae.enable_slicing()
2121

22-
def get_seed_generators(amount, seed):
23-
generator = [torch.Generator(device="cuda").manual_seed(seed + i) for i in range(amount)]
24-
return generator
25-
26-
2722
def generate_auraflow(prompt,
2823
width,
2924
height,
@@ -32,7 +27,8 @@ def generate_auraflow(prompt,
3227
negative_prompt=None,
3328
guidance_scale=None,
3429
seed=None,
35-
model_name="fal/AuraFlow"):
30+
model_name="fal/AuraFlow",
31+
lora_name=None):
3632
global PIPELINE
3733
global LOADED
3834
if not LOADED:
@@ -48,6 +44,8 @@ def generate_auraflow(prompt,
4844
kwargs["negative_prompt"] = negative_prompt
4945
if seed is not None:
5046
kwargs["generator"] = get_seed_generators(kwargs["num_images_per_prompt"], seed)
47+
if lora_name is not None:
48+
PIPELINE = load_loras(PIPELINE, "auraflow", lora_name)
5149
try:
5250
images = PIPELINE(**kwargs).images
5351
return {"status": True,
@@ -72,6 +70,8 @@ def auraflow_generate(data: AuraFlowRequest = Body(...)):
7270
kwargs["guidance_scale"] = data.guidance_scale
7371
if data.seed:
7472
kwargs["seed"] = data.seed
73+
if data.lora_name:
74+
kwargs["lora_name"] = data.lora_name
7575
try:
7676
response = generate_auraflow(**kwargs)
7777
if response["status"] is True:

‎modules/avernus_client.py‎

Lines changed: 31 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,7 @@ async def ace_music(self, prompt, lyrics, audio_duration=None, guidance_scale=No
3939
return {"ERROR": str(e)}
4040

4141
async def auraflow_image(self, prompt, negative_prompt=None, model_name=None, width=None, height=None, steps=None,
42-
batch_size=None, seed=None, guidance_scale=None):
42+
batch_size=None, seed=None, guidance_scale=None, lora_name=None):
4343
"""This takes a prompt and optional other variables and returns a list of base64 encoded images"""
4444
url = f"http://{self.base_url}/auraflow_generate"
4545
data = {"prompt": prompt,
@@ -50,7 +50,8 @@ async def auraflow_image(self, prompt, negative_prompt=None, model_name=None, wi
5050
"steps": steps,
5151
"batch_size": batch_size,
5252
"seed": seed,
53-
"guidance_scale": guidance_scale}
53+
"guidance_scale": guidance_scale,
54+
"lora_name": lora_name}
5455
try:
5556
async with httpx.AsyncClient(timeout=None) as client:
5657
response = await client.post(url, json=data)
@@ -257,7 +258,7 @@ async def flux_kontext(self, prompt, negative_prompt=None, image=None, model_nam
257258
return {"ERROR": str(e)}
258259

259260
async def framepack(self, prompt, image, negative_prompt=None, width=None, height=None, steps=None, num_frames=None,
260-
guidance_scale=None, last_image=None, seed=None, model_name=None):
261+
guidance_scale=None, last_image=None, seed=None, model_name=None, lora_name=None):
261262
"""This takes a prompt and returns a video"""
262263
url = f"http://{self.base_url}/framepack_generate"
263264
data = {"prompt": prompt,
@@ -270,7 +271,8 @@ async def framepack(self, prompt, image, negative_prompt=None, width=None, heigh
270271
"steps": steps,
271272
"image": image,
272273
"last_image": last_image,
273-
"model_name": model_name}
274+
"model_name": model_name,
275+
"lora_name": lora_name}
274276
try:
275277
async with httpx.AsyncClient(timeout=None) as client:
276278
response = await client.post(url, json=data)
@@ -286,7 +288,7 @@ async def framepack(self, prompt, image, negative_prompt=None, width=None, heigh
286288
return {"ERROR": str(e)}
287289

288290
async def hidream_image(self, prompt, negative_prompt=None, model_name=None, width=None, height=None, steps=None,
289-
batch_size=None, seed=None, guidance_scale=None):
291+
batch_size=None, seed=None, guidance_scale=None, lora_name=None):
290292
"""This takes a prompt and optional other variables and returns a list of base64 encoded images"""
291293
url = f"http://{self.base_url}/hidream_generate"
292294
data = {"prompt": prompt,
@@ -297,7 +299,8 @@ async def hidream_image(self, prompt, negative_prompt=None, model_name=None, wid
297299
"steps": steps,
298300
"batch_size": batch_size,
299301
"seed": seed,
300-
"guidance_scale": guidance_scale}
302+
"guidance_scale": guidance_scale,
303+
"lora_name": lora_name}
301304
try:
302305
async with httpx.AsyncClient(timeout=None) as client:
303306
response = await client.post(url, json=data)
@@ -310,7 +313,7 @@ async def hidream_image(self, prompt, negative_prompt=None, model_name=None, wid
310313
return {"ERROR": str(e)}
311314

312315
async def hunyuan_ti2v(self, prompt, negative_prompt=None, width=None, height=None, steps=None, num_frames=None,
313-
guidance_scale=None, image=None, seed=None, model_name=None, flow_shift=None):
316+
guidance_scale=None, image=None, seed=None, model_name=None, flow_shift=None, lora_name=None):
314317
"""This takes a prompt and returns a video"""
315318
url = f"http://{self.base_url}/hunyuan_ti2v_generate"
316319
data = {"prompt": prompt,
@@ -323,7 +326,8 @@ async def hunyuan_ti2v(self, prompt, negative_prompt=None, width=None, height=No
323326
"steps": steps,
324327
"image": image,
325328
"model_name": model_name,
326-
"flow_shift": flow_shift}
329+
"flow_shift": flow_shift,
330+
"lora_name": lora_name}
327331
try:
328332
async with httpx.AsyncClient(timeout=None) as client:
329333
response = await client.post(url, json=data)
@@ -360,7 +364,7 @@ async def image_gen_aux_upscale(self, image, model=None, scale=None, tiling=None
360364
return {"ERROR": str(e)}
361365

362366
async def kandinsky5_t2v(self, prompt, negative_prompt=None, width=None, height=None, steps=None, num_frames=None,
363-
guidance_scale=None, seed=None, model_name=None):
367+
guidance_scale=None, seed=None, model_name=None, lora_name=None):
364368
"""This takes a prompt and returns a video"""
365369
url = f"http://{self.base_url}/kandinsky5_t2v_generate"
366370
data = {"prompt": prompt,
@@ -371,7 +375,8 @@ async def kandinsky5_t2v(self, prompt, negative_prompt=None, width=None, height=
371375
"guidance_scale": guidance_scale,
372376
"seed": seed,
373377
"steps": steps,
374-
"model_name": model_name}
378+
"model_name": model_name,
379+
"lora_name": lora_name}
375380
try:
376381
async with httpx.AsyncClient(timeout=None) as client:
377382
response = await client.post(url, json=data)
@@ -517,7 +522,7 @@ async def llm_chat(self, prompt, model_name=None, messages=None):
517522
return {"ERROR": str(e)}
518523

519524
async def ltx_ti2v(self, prompt, negative_prompt=None, width=None, height=None, steps=None, num_frames=None,
520-
guidance_scale=None, image=None, seed=None, model_name=None, frame_rate=None):
525+
guidance_scale=None, image=None, seed=None, model_name=None, frame_rate=None, lora_name=None):
521526
"""This takes a prompt and returns a video"""
522527
url = f"http://{self.base_url}/ltx_ti2v_generate"
523528
data = {"prompt": prompt,
@@ -530,7 +535,8 @@ async def ltx_ti2v(self, prompt, negative_prompt=None, width=None, height=None,
530535
"steps": steps,
531536
"image": image,
532537
"model_name": model_name,
533-
"frame_rate": frame_rate}
538+
"frame_rate": frame_rate,
539+
"lora_name": lora_name}
534540
try:
535541
async with httpx.AsyncClient(timeout=None) as client:
536542
response = await client.post(url, json=data)
@@ -546,7 +552,7 @@ async def ltx_ti2v(self, prompt, negative_prompt=None, width=None, height=None,
546552
return {"ERROR": str(e)}
547553

548554
async def lumina2_image(self, prompt, negative_prompt=None, model_name=None, width=None, height=None, steps=None,
549-
batch_size=None, seed=None, guidance_scale=None):
555+
batch_size=None, seed=None, guidance_scale=None, lora_name=None):
550556
"""This takes a prompt and optional other variables and returns a list of base64 encoded images"""
551557
url = f"http://{self.base_url}/lumina2_generate"
552558
data = {"prompt": prompt,
@@ -557,7 +563,8 @@ async def lumina2_image(self, prompt, negative_prompt=None, model_name=None, wid
557563
"steps": steps,
558564
"batch_size": batch_size,
559565
"seed": seed,
560-
"guidance_scale": guidance_scale}
566+
"guidance_scale": guidance_scale,
567+
"lora_name": lora_name}
561568
try:
562569
async with httpx.AsyncClient(timeout=None) as client:
563570
response = await client.post(url, json=data)
@@ -974,7 +981,7 @@ async def swin2sr(self, image):
974981
return {"ERROR": str(e)}
975982

976983
async def wan_ti2v(self, prompt, negative_prompt=None, width=None, height=None, steps=None, num_frames=None,
977-
guidance_scale=None, image=None, seed=None, model_name=None, flow_shift=None):
984+
guidance_scale=None, image=None, seed=None, model_name=None, flow_shift=None, lora_name=None):
978985
"""This takes a prompt and returns a video"""
979986
url = f"http://{self.base_url}/wan_ti2v_generate"
980987
data = {"prompt": prompt,
@@ -987,7 +994,8 @@ async def wan_ti2v(self, prompt, negative_prompt=None, width=None, height=None,
987994
"steps": steps,
988995
"image": image,
989996
"model_name": model_name,
990-
"flow_shift": flow_shift}
997+
"flow_shift": flow_shift,
998+
"lora_name": lora_name}
991999
try:
9921000
async with httpx.AsyncClient(timeout=None) as client:
9931001
response = await client.post(url, json=data)
@@ -1003,7 +1011,8 @@ async def wan_ti2v(self, prompt, negative_prompt=None, width=None, height=None,
10031011
return {"ERROR": str(e)}
10041012

10051013
async def wan_vace(self, prompt, negative_prompt=None, width=None, height=None, steps=None, num_frames=None,
1006-
guidance_scale=None, first_frame=None, last_frame=None, flow_shift=None, seed=None, model_name=None):
1014+
guidance_scale=None, first_frame=None, last_frame=None, flow_shift=None, seed=None,
1015+
model_name=None, lora_name=None):
10071016
"""This takes a prompt and returns a video"""
10081017
url = f"http://{self.base_url}/wan_vace_generate"
10091018
data = {"prompt": prompt,
@@ -1017,7 +1026,8 @@ async def wan_vace(self, prompt, negative_prompt=None, width=None, height=None,
10171026
"first_frame": first_frame,
10181027
"last_frame": last_frame,
10191028
"flow_shift": flow_shift,
1020-
"model_name": model_name}
1029+
"model_name": model_name,
1030+
"lora_name": lora_name}
10211031
try:
10221032
async with httpx.AsyncClient(timeout=None) as client:
10231033
response = await client.post(url, json=data)
@@ -1033,7 +1043,7 @@ async def wan_vace(self, prompt, negative_prompt=None, width=None, height=None,
10331043
return {"ERROR": str(e)}
10341044

10351045
async def wan_v2v(self, prompt, negative_prompt=None, width=None, height=None, steps=None,
1036-
guidance_scale=None, seed=None, model_name=None, video_path=None, flow_shift=None):
1046+
guidance_scale=None, seed=None, model_name=None, video_path=None, flow_shift=None, lora_name=None):
10371047
"""This takes a prompt and (optionally) a video, and returns a generated video."""
10381048
url = f"http://{self.base_url}/wan_v2v_generate"
10391049
data = {
@@ -1045,8 +1055,8 @@ async def wan_v2v(self, prompt, negative_prompt=None, width=None, height=None, s
10451055
"seed": seed,
10461056
"steps": steps,
10471057
"model_name": model_name,
1048-
"flow_shift": flow_shift
1049-
}
1058+
"flow_shift": flow_shift,
1059+
"lora_name": lora_name}
10501060
data = {k: str(v) for k, v in data.items() if v is not None}
10511061
files = None
10521062
if video_path:

‎modules/chroma.py‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
import torch
66

77
from pydantic_models import ChromaRequest, ImageResponse
8-
from utils import base64_to_image, image_to_base64
8+
from utils import image_to_base64, get_seed_generators, load_loras
99

1010
PIPELINE: ChromaPipeline
1111
LOADED: bool = False
@@ -19,11 +19,6 @@ def load_chroma_pipeline(model_name="Meatfucker/Chroma1-HD-bnb-int8"):
1919
PIPELINE.enable_model_cpu_offload()
2020
PIPELINE.vae.enable_tiling()
2121

22-
def get_seed_generators(amount, seed):
23-
generator = [torch.Generator(device="cuda").manual_seed(seed + i) for i in range(amount)]
24-
return generator
25-
26-
2722
def generate_chroma(prompt,
2823
width,
2924
height,
@@ -32,7 +27,8 @@ def generate_chroma(prompt,
3227
negative_prompt=None,
3328
guidance_scale=None,
3429
seed=None,
35-
model_name=None):
30+
model_name=None,
31+
lora_name=None):
3632
global PIPELINE
3733
global LOADED
3834
if not LOADED:
@@ -50,6 +46,8 @@ def generate_chroma(prompt,
5046
"guidance_scale": guidance_scale if guidance_scale is not None else 5.0}
5147
if seed is not None:
5248
kwargs["generator"] = get_seed_generators(kwargs["num_images_per_prompt"], seed)
49+
if lora_name is not None:
50+
PIPELINE = load_loras(PIPELINE, "chroma", lora_name)
5351
try:
5452
images = PIPELINE(**kwargs).images
5553
return {"status": True,
@@ -74,6 +72,8 @@ def chroma_generate(data: ChromaRequest = Body(...)):
7472
kwargs["guidance_scale"] = data.guidance_scale
7573
if data.seed:
7674
kwargs["seed"] = data.seed
75+
if data.lora_name:
76+
kwargs["lora_name"] = data.lora_name
7777
try:
7878
response = generate_chroma(**kwargs)
7979
if response["status"] is True:

‎modules/chroma_i2i.py‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
import torch
66

77
from pydantic_models import ChromaRequest, ImageResponse
8-
from utils import base64_to_image, image_to_base64
8+
from utils import base64_to_image, image_to_base64, get_seed_generators, load_loras
99

1010
PIPELINE: ChromaImg2ImgPipeline
1111
LOADED: bool = False
@@ -19,11 +19,6 @@ def load_chroma_pipeline(model_name="Meatfucker/Chroma1-HD-bnb-int8"):
1919
PIPELINE.enable_model_cpu_offload()
2020
PIPELINE.vae.enable_tiling()
2121

22-
def get_seed_generators(amount, seed):
23-
generator = [torch.Generator(device="cuda").manual_seed(seed + i) for i in range(amount)]
24-
return generator
25-
26-
2722
def generate_chroma(prompt,
2823
image,
2924
width,
@@ -34,7 +29,8 @@ def generate_chroma(prompt,
3429
guidance_scale=None,
3530
seed=None,
3631
model_name=None,
37-
strength=None):
32+
strength=None,
33+
lora_name=None):
3834
global PIPELINE
3935
global LOADED
4036
if not LOADED:
@@ -54,6 +50,8 @@ def generate_chroma(prompt,
5450
"strength": strength if strength is not None else 0.9}
5551
if seed is not None:
5652
kwargs["generator"] = get_seed_generators(kwargs["num_images_per_prompt"], seed)
53+
if lora_name is not None:
54+
PIPELINE = load_loras(PIPELINE, "chroma", lora_name)
5755
try:
5856
images = PIPELINE(**kwargs).images
5957
return {"status": True,
@@ -82,6 +80,8 @@ def chroma_generate(data: ChromaRequest = Body(...)):
8280
kwargs["seed"] = data.seed
8381
if data.strength:
8482
kwargs["strength"] = data.strength
83+
if data.lora_name:
84+
kwargs["lora_name"] = data.lora_name
8585
try:
8686
response = generate_chroma(**kwargs)
8787
if response["status"] is True:

‎modules/chronoedit.py‎

Lines changed: 2 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010

1111
from ChronoEdit.chronoedit_diffusers.pipeline_chronoedit import ChronoEditPipeline
1212
from pydantic_models import ChronoEditRequest, ImageResponse
13-
from utils import base64_to_image, image_to_base64
13+
from utils import base64_to_image, image_to_base64, get_seed_generators, load_loras
1414

1515
PIPELINE: ChronoEditPipeline
1616
LOADED: bool = False
@@ -44,25 +44,6 @@ def load_chronoedit_pipeline(model_name="Meatfucker/ChronoEdit-bnb-nf4", flow_sh
4444
PIPELINE.scheduler = UniPCMultistepScheduler.from_config(PIPELINE.scheduler.config, flow_shift=flow_shift)
4545
PIPELINE.enable_model_cpu_offload()
4646

47-
def get_seed_generators(amount, seed):
48-
generator = [torch.Generator(device="cuda").manual_seed(seed + i) for i in range(amount)]
49-
return generator
50-
51-
def load_chronoedit_loras(lora_name):
52-
global PIPELINE
53-
try:
54-
lora_list = []
55-
for lora in lora_name:
56-
try:
57-
lora_name = os.path.splitext(lora)[0]
58-
PIPELINE.load_lora_weights(f"loras/chronoedit/{lora}", adapter_name=lora_name)
59-
lora_list.append(lora_name)
60-
except Exception:
61-
pass
62-
PIPELINE.set_adapters(lora_list)
63-
except Exception:
64-
pass
65-
6647
def generate_chronoedit_image(prompt,
6748
width,
6849
height,
@@ -98,9 +79,8 @@ def generate_chronoedit_image(prompt,
9879
kwargs["image"] = image
9980
kwargs["width"] = width
10081
kwargs["height"] = height
101-
10282
if lora_name is not None:
103-
load_chronoedit_loras(lora_name)
83+
PIPELINE = load_loras(PIPELINE, "chronoedit", lora_name)
10484

10585
frames = PIPELINE(**kwargs).frames
10686
images = []

0 commit comments

Comments
 (0)