Update backend.py

flux_app/backend.py

@@ -1,4 +1,3 @@
-# backend.py
 import torch
 from diffusers import (
     DiffusionPipeline,
@@ -11,6 +10,10 @@ from flux_app.utilities import calculate_shift, retrieve_timesteps, load_image_f
 from flux_app.lora_handling import flux_pipe_call_that_returns_an_iterable_of_images  # Absolute import
 import time
 from huggingface_hub import login
+import spaces
+# Ensure CUDA is available
+if not torch.cuda.is_available():
+    raise RuntimeError("CUDA is not available. Please run on a GPU-enabled environment.")
 
 class ModelManager:
     def __init__(self, hf_token=None):
@@ -18,12 +21,16 @@ class ModelManager:
         self.pipe_i2i = None
         self.good_vae = None
         self.taef1 = None
+        
+        # Clear CUDA memory cache before loading models
+        torch.cuda.empty_cache()
 
         if hf_token:
             login(token=hf_token)  # Log in with the provided token
 
         self.initialize_models()
-
+    
+    @spaces.GPU(duration=300)
     def initialize_models(self):
         """Initializes the diffusion pipelines and autoencoders."""
         self.taef1 = AutoencoderTiny.from_pretrained(TAEF1_MODEL, torch_dtype=DTYPE).to(DEVICE)
@@ -38,17 +45,18 @@ class ModelManager:
             text_encoder_2=self.pipe.text_encoder_2,
             tokenizer_2=self.pipe.tokenizer_2,
             torch_dtype=DTYPE
-        )
+        ).to(DEVICE)
 
         setattr(self.pipe, "flux_pipe_call_that_returns_an_iterable_of_images", self.process_images)
-
+    @spaces.GPU(duration=300)
     def process_images(self, *args, **kwargs):
         return flux_pipe_call_that_returns_an_iterable_of_images(self.pipe, *args, **kwargs)
-
+    @spaces.GPU(duration=300)
     def generate_image(self, prompt_mash, steps, seed, cfg_scale, width, height, lora_scale):
-        """Generates an image using the text-to-image pipeline."""
-        self.pipe.to(DEVICE)
+        """Generates an image using the FLUX pipeline."""
+        self.pipe.to(DEVICE)  # Ensure pipeline is on GPU
         generator = torch.Generator(device=DEVICE).manual_seed(seed)
+        
         with calculateDuration("Generating image"):
             for img in self.pipe.flux_pipe_call_that_returns_an_iterable_of_images(
                 prompt=prompt_mash,
@@ -62,23 +70,28 @@ class ModelManager:
                 good_vae=self.good_vae,
             ):
                 yield img
-
+    @spaces.GPU(duration=300)
     def generate_image_to_image(self, prompt_mash, image_input_path, image_strength, steps, cfg_scale, width, height, lora_scale, seed):
-        """Generates an image using the image-to-image pipeline."""
+        """Generates an image using image-to-image processing."""
         generator = torch.Generator(device=DEVICE).manual_seed(seed)
         self.pipe_i2i.to(DEVICE)
         image_input = load_image_from_path(image_input_path)
-        with calculateDuration("Generating image to image"):
-            final_image = self.pipe_i2i(
-                prompt=prompt_mash,
-                image=image_input,
-                strength=image_strength,
-                num_inference_steps=steps,
-                guidance_scale=cfg_scale,
-                width=width,
-                height=height,
-                generator=generator,
-                joint_attention_kwargs={"scale": lora_scale},
-                output_type="pil",
-            ).images[0]
-            return final_image
+        
+        final_image = self.pipe_i2i(
+            prompt=prompt_mash,
+            image=image_input,
+            strength=image_strength,
+            num_inference_steps=steps,
+            guidance_scale=cfg_scale,
+            width=width,
+            height=height,
+            generator=generator,
+            joint_attention_kwargs={"scale": lora_scale},
+            output_type="pil",
+        ).images[0]
+        return final_image
+
+# Ensure the pipeline is properly initialized when running
+if __name__ == "__main__":
+    model_manager = ModelManager()
+    print("Model Manager initialized successfully with GPU support.")