ruslanmv commited on
Commit
b841fb8
·
1 Parent(s): b1d31cc

Update backend.py

Browse files
Files changed (1) hide show
  1. 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 text-to-image pipeline."""
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 the image-to-image pipeline."""
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
- with calculateDuration("Generating image to image"):
72
- final_image = self.pipe_i2i(
73
- prompt=prompt_mash,
74
- image=image_input,
75
- strength=image_strength,
76
- num_inference_steps=steps,
77
- guidance_scale=cfg_scale,
78
- width=width,
79
- height=height,
80
- generator=generator,
81
- joint_attention_kwargs={"scale": lora_scale},
82
- output_type="pil",
83
- ).images[0]
84
- return final_image
 
 
 
 
 
 
 
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.")