Spaces:
Running
on
Zero
Running
on
Zero
Update backend.py
Browse files- flux_app/backend.py +36 -23
flux_app/backend.py
CHANGED
@@ -1,4 +1,3 @@
|
|
1 |
-
# backend.py
|
2 |
import torch
|
3 |
from diffusers import (
|
4 |
DiffusionPipeline,
|
@@ -11,6 +10,10 @@ from flux_app.utilities import calculate_shift, retrieve_timesteps, load_image_f
|
|
11 |
from flux_app.lora_handling import flux_pipe_call_that_returns_an_iterable_of_images # Absolute import
|
12 |
import time
|
13 |
from huggingface_hub import login
|
|
|
|
|
|
|
|
|
14 |
|
15 |
class ModelManager:
|
16 |
def __init__(self, hf_token=None):
|
@@ -18,12 +21,16 @@ class ModelManager:
|
|
18 |
self.pipe_i2i = None
|
19 |
self.good_vae = None
|
20 |
self.taef1 = None
|
|
|
|
|
|
|
21 |
|
22 |
if hf_token:
|
23 |
login(token=hf_token) # Log in with the provided token
|
24 |
|
25 |
self.initialize_models()
|
26 |
-
|
|
|
27 |
def initialize_models(self):
|
28 |
"""Initializes the diffusion pipelines and autoencoders."""
|
29 |
self.taef1 = AutoencoderTiny.from_pretrained(TAEF1_MODEL, torch_dtype=DTYPE).to(DEVICE)
|
@@ -38,17 +45,18 @@ class ModelManager:
|
|
38 |
text_encoder_2=self.pipe.text_encoder_2,
|
39 |
tokenizer_2=self.pipe.tokenizer_2,
|
40 |
torch_dtype=DTYPE
|
41 |
-
)
|
42 |
|
43 |
setattr(self.pipe, "flux_pipe_call_that_returns_an_iterable_of_images", self.process_images)
|
44 |
-
|
45 |
def process_images(self, *args, **kwargs):
|
46 |
return flux_pipe_call_that_returns_an_iterable_of_images(self.pipe, *args, **kwargs)
|
47 |
-
|
48 |
def generate_image(self, prompt_mash, steps, seed, cfg_scale, width, height, lora_scale):
|
49 |
-
"""Generates an image using the
|
50 |
-
self.pipe.to(DEVICE)
|
51 |
generator = torch.Generator(device=DEVICE).manual_seed(seed)
|
|
|
52 |
with calculateDuration("Generating image"):
|
53 |
for img in self.pipe.flux_pipe_call_that_returns_an_iterable_of_images(
|
54 |
prompt=prompt_mash,
|
@@ -62,23 +70,28 @@ class ModelManager:
|
|
62 |
good_vae=self.good_vae,
|
63 |
):
|
64 |
yield img
|
65 |
-
|
66 |
def generate_image_to_image(self, prompt_mash, image_input_path, image_strength, steps, cfg_scale, width, height, lora_scale, seed):
|
67 |
-
"""Generates an image using
|
68 |
generator = torch.Generator(device=DEVICE).manual_seed(seed)
|
69 |
self.pipe_i2i.to(DEVICE)
|
70 |
image_input = load_image_from_path(image_input_path)
|
71 |
-
|
72 |
-
|
73 |
-
|
74 |
-
|
75 |
-
|
76 |
-
|
77 |
-
|
78 |
-
|
79 |
-
|
80 |
-
|
81 |
-
|
82 |
-
|
83 |
-
|
84 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
import torch
|
2 |
from diffusers import (
|
3 |
DiffusionPipeline,
|
|
|
10 |
from flux_app.lora_handling import flux_pipe_call_that_returns_an_iterable_of_images # Absolute import
|
11 |
import time
|
12 |
from huggingface_hub import login
|
13 |
+
import spaces
|
14 |
+
# Ensure CUDA is available
|
15 |
+
if not torch.cuda.is_available():
|
16 |
+
raise RuntimeError("CUDA is not available. Please run on a GPU-enabled environment.")
|
17 |
|
18 |
class ModelManager:
|
19 |
def __init__(self, hf_token=None):
|
|
|
21 |
self.pipe_i2i = None
|
22 |
self.good_vae = None
|
23 |
self.taef1 = None
|
24 |
+
|
25 |
+
# Clear CUDA memory cache before loading models
|
26 |
+
torch.cuda.empty_cache()
|
27 |
|
28 |
if hf_token:
|
29 |
login(token=hf_token) # Log in with the provided token
|
30 |
|
31 |
self.initialize_models()
|
32 |
+
|
33 |
+
@spaces.GPU(duration=300)
|
34 |
def initialize_models(self):
|
35 |
"""Initializes the diffusion pipelines and autoencoders."""
|
36 |
self.taef1 = AutoencoderTiny.from_pretrained(TAEF1_MODEL, torch_dtype=DTYPE).to(DEVICE)
|
|
|
45 |
text_encoder_2=self.pipe.text_encoder_2,
|
46 |
tokenizer_2=self.pipe.tokenizer_2,
|
47 |
torch_dtype=DTYPE
|
48 |
+
).to(DEVICE)
|
49 |
|
50 |
setattr(self.pipe, "flux_pipe_call_that_returns_an_iterable_of_images", self.process_images)
|
51 |
+
@spaces.GPU(duration=300)
|
52 |
def process_images(self, *args, **kwargs):
|
53 |
return flux_pipe_call_that_returns_an_iterable_of_images(self.pipe, *args, **kwargs)
|
54 |
+
@spaces.GPU(duration=300)
|
55 |
def generate_image(self, prompt_mash, steps, seed, cfg_scale, width, height, lora_scale):
|
56 |
+
"""Generates an image using the FLUX pipeline."""
|
57 |
+
self.pipe.to(DEVICE) # Ensure pipeline is on GPU
|
58 |
generator = torch.Generator(device=DEVICE).manual_seed(seed)
|
59 |
+
|
60 |
with calculateDuration("Generating image"):
|
61 |
for img in self.pipe.flux_pipe_call_that_returns_an_iterable_of_images(
|
62 |
prompt=prompt_mash,
|
|
|
70 |
good_vae=self.good_vae,
|
71 |
):
|
72 |
yield img
|
73 |
+
@spaces.GPU(duration=300)
|
74 |
def generate_image_to_image(self, prompt_mash, image_input_path, image_strength, steps, cfg_scale, width, height, lora_scale, seed):
|
75 |
+
"""Generates an image using image-to-image processing."""
|
76 |
generator = torch.Generator(device=DEVICE).manual_seed(seed)
|
77 |
self.pipe_i2i.to(DEVICE)
|
78 |
image_input = load_image_from_path(image_input_path)
|
79 |
+
|
80 |
+
final_image = self.pipe_i2i(
|
81 |
+
prompt=prompt_mash,
|
82 |
+
image=image_input,
|
83 |
+
strength=image_strength,
|
84 |
+
num_inference_steps=steps,
|
85 |
+
guidance_scale=cfg_scale,
|
86 |
+
width=width,
|
87 |
+
height=height,
|
88 |
+
generator=generator,
|
89 |
+
joint_attention_kwargs={"scale": lora_scale},
|
90 |
+
output_type="pil",
|
91 |
+
).images[0]
|
92 |
+
return final_image
|
93 |
+
|
94 |
+
# Ensure the pipeline is properly initialized when running
|
95 |
+
if __name__ == "__main__":
|
96 |
+
model_manager = ModelManager()
|
97 |
+
print("Model Manager initialized successfully with GPU support.")
|