Drastic Update to Image Generation to accomodate @spaces.GPU() excessive time usage

app.py

@@ -8,6 +8,7 @@ from typing import Optional, Union, List, Tuple
 
 from PIL import Image, ImageFilter
 import cv2
+
 import utils.constants as constants
 
 from haishoku.haishoku import Haishoku
@@ -91,6 +92,7 @@ from utils.version_info import (
     #release_torch_resources,
     #get_torch_info
 )
+from src.condition import Condition
 import spaces
 
 input_image_palette = []
@@ -199,113 +201,24 @@ condition_dict = {
     "fill": 9,
 }
 
-class Condition(object):
-    def __init__(
-        self,
-        condition_type: str,
-        raw_img: Union[Image.Image, torch.Tensor] = None,
-        condition: Union[Image.Image,torch.Tensor] = None,
-        mask=None,
-    ) -> None:
-        self.condition_type = condition_type
-        assert raw_img is not None or condition is not None
-        if raw_img is not None:
-            self.condition = self.get_condition(condition_type, raw_img)
-        else:
-            self.condition = condition
-        # TODO: Add mask support
-        assert mask is None, "Mask not supported yet"
-    def get_condition(
-        self, condition_type: str, raw_img: Union[Image.Image, torch.Tensor]
-    ) -> Union[Image.Image, torch.Tensor]:
-        """
-        Returns the condition image.
-        """
-        if condition_type == "depth":
-            from transformers import pipeline
-            depth_pipe = pipeline(
-                task="depth-estimation",
-                model="LiheYoung/depth-anything-small-hf",
-                device="cuda",
-            )
-            source_image = raw_img.convert("RGB")
-            condition_img = depth_pipe(source_image)["depth"].convert("RGB")
-            return condition_img
-        elif condition_type == "canny":
-            img = np.array(raw_img)
-            edges = cv2.Canny(img, 100, 200)
-            edges = Image.fromarray(edges).convert("RGB")
-            return edges
-        elif condition_type == "subject":
-            return raw_img
-        elif condition_type == "coloring":
-            return raw_img.convert("L").convert("RGB")
-        elif condition_type == "deblurring":
-            condition_image = (
-                raw_img.convert("RGB")
-                .filter(ImageFilter.GaussianBlur(10))
-                .convert("RGB")
-            )
-            return condition_image
-        elif condition_type == "fill":
-            return raw_img.convert("RGB")
-        return self.condition
-    @property
-    def type_id(self) -> int:
-        """
-        Returns the type id of the condition.
-        """
-        return condition_dict[self.condition_type]
-    @classmethod
-    def get_type_id(cls, condition_type: str) -> int:
-        """
-        Returns the type id of the condition.
-        """
-        return condition_dict[condition_type]
-    def _encode_image(self, pipe: FluxPipeline, cond_img: Image.Image) -> torch.Tensor:
-        """
-        Encodes an image condition into tokens using the pipeline.
-        """
-        cond_img = pipe.image_processor.preprocess(cond_img)
-        cond_img = cond_img.to(pipe.device).to(pipe.dtype)
-        cond_img = pipe.vae.encode(cond_img).latent_dist.sample()
-        cond_img = (
-            cond_img - pipe.vae.config.shift_factor
-        ) * pipe.vae.config.scaling_factor
-        cond_tokens = pipe._pack_latents(cond_img, *cond_img.shape)
-        cond_ids = pipe._prepare_latent_image_ids(
-            cond_img.shape[0],
-            cond_img.shape[2]//2,
-            cond_img.shape[3]//2,
-            pipe.device,
-            pipe.dtype,
-        )
-        return cond_tokens, cond_ids
-    def encode(self, pipe: FluxPipeline) -> Tuple[torch.Tensor, torch.Tensor, int]:
-        """
-        Encodes the condition into tokens, ids and type_id.
-        """
-        if self.condition_type in [
-            "depth",
-            "canny",
-            "subject",
-            "coloring",
-            "deblurring",
-            "fill",
-        ]:
-            tokens, ids = self._encode_image(pipe, self.condition)
-        else:
-            raise NotImplementedError(
-                f"Condition type {self.condition_type} not implemented"
-            )
-        type_id = torch.ones_like(ids[:, :1]) * self.type_id
-        return tokens, ids, type_id
+@spaces.GPU(duration=120, progress=gr.Progress(track_tqdm=True))
+def generate_image(pipe, conditions, generate_params, progress=gr.Progress(track_tqdm=True)):
+    gr.Info("Generating AI image...",duration=5)
+    result = pipe(**generate_params) 
+    image = result.images[0]
+    # Clean up
+    del result
+    del conditions   
+    # Delete the pipeline and clear cache
+    del pipe
+    torch.cuda.empty_cache()
+    torch.cuda.ipc_collect()
+    print(torch.cuda.memory_summary(device=None, abbreviated=False))
+    return image
 
-# @spaces.GPU(duration=140, progress=gr.Progress(track_tqdm=True))
-# def generate_image(pipe, generate_params, progress=gr.Progress(track_tqdm=True)):
-#     return pipe(**generate_params)
 
-@spaces.GPU(duration=200, progress=gr.Progress(track_tqdm=True))
+@spaces.GPU(duration=90)
+@torch.no_grad()
 def generate_image_lowmem(
     text,
     neg_prompt=None,
@@ -331,195 +244,205 @@ def generate_image_lowmem(
                         f"Available options: {list(PIPELINE_CLASSES.keys())}")
 
     #initialize_cuda()
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    #from src.condition import Condition
+    device = "cuda" if torch.cuda.is_available() else "cpu"    
 
     print(f"device:{device}\nmodel_name:{model_name}\nlora_weights:{lora_weights}\n")
     #print(f"\n {get_torch_info()}\n")
     # Disable gradient calculations
-    with torch.no_grad():
-        # Initialize the pipeline inside the context manager
-        pipe = pipeline_class.from_pretrained(
-            model_name,
-            torch_dtype=torch.bfloat16 if device == "cuda" else torch.float32
-        ).to(device)
-        # Optionally, don't use CPU offload if not necessary
+#with torch.no_grad():
+    gr.Info("Initialize the pipeline inside the context manager",duration=5)
+    # Initialize the pipeline inside the context manager
+    pipe = pipeline_class.from_pretrained(
+        model_name,
+        torch_dtype=torch.bfloat16 if device == "cuda" else torch.float32
+    ).to(device)
+    # Optionally, don't use CPU offload if not necessary
         
-        # alternative version that may be more efficient
-        # pipe.enable_sequential_cpu_offload()
-        if pipeline_name == "FluxPipeline":
-            pipe.enable_model_cpu_offload()
-            pipe.vae.enable_slicing()
-            #pipe.vae.enable_tiling()
-        else:
-            pipe.enable_model_cpu_offload()
-
-        # Access the tokenizer from the pipeline
-        tokenizer = pipe.tokenizer
-
-        # Check if add_prefix_space is set and convert to slow tokenizer if necessary
-        if getattr(tokenizer, 'add_prefix_space', False):
-            tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False, device_map = 'cpu')
-            # Update the pipeline's tokenizer
-            pipe.tokenizer = tokenizer
-            pipe.to(device)
-
-        flash_attention_enabled = torch.backends.cuda.flash_sdp_enabled()
-        if flash_attention_enabled == False:
-            #Enable xFormers memory-efficient attention (optional)
-            #pipe.enable_xformers_memory_efficient_attention()
-            print("\nEnabled xFormers memory-efficient attention.\n")
-        else:            
-            pipe.attn_implementation="flash_attention_2"
-            print("\nEnabled flash_attention_2.\n")
-
-        condition_type = "subject"
-        # Load LoRA weights
-        # note: does not yet handle multiple LoRA weights with different names, needs .set_adapters(["depth", "hyper-sd"], adapter_weights=[0.85, 0.125])        
-        if lora_weights:
-            for lora_weight in lora_weights:
-                lora_configs = constants.LORA_DETAILS.get(lora_weight, [])
-                lora_weight_set = False
-                if lora_configs:
-                    for config in lora_configs:
-                        # Load LoRA weights with optional weight_name and adapter_name
-                        if 'weight_name' in config:
-                            weight_name = config.get("weight_name")
-                            adapter_name = config.get("adapter_name")
-                            lora_collection = config.get("lora_collection")
-                            if weight_name and adapter_name and lora_collection and lora_weight_set == False:
-                                pipe.load_lora_weights(
-                                    lora_collection,
-                                    weight_name=weight_name,
-                                    adapter_name=adapter_name,
-                                    token=constants.HF_API_TOKEN
-                                )
-                                lora_weight_set = True
-                                print(f"\npipe.load_lora_weights({lora_weight}, weight_name={weight_name}, adapter_name={adapter_name}, lora_collection={lora_collection}\n")  
-                            elif weight_name and adapter_name==None and lora_collection and lora_weight_set == False:
-                                pipe.load_lora_weights(
-                                    lora_collection,
-                                    weight_name=weight_name,
-                                    token=constants.HF_API_TOKEN
-                                )
-                                lora_weight_set = True
-                                print(f"\npipe.load_lora_weights({lora_weight}, weight_name={weight_name}, adapter_name={adapter_name}, lora_collection={lora_collection}\n")  
-                            elif weight_name and adapter_name and lora_weight_set == False:
-                                pipe.load_lora_weights(
-                                    lora_weight,
-                                    weight_name=weight_name,
-                                    adapter_name=adapter_name,
-                                    token=constants.HF_API_TOKEN
-                                )
-                                lora_weight_set = True
-                                print(f"\npipe.load_lora_weights({lora_weight}, weight_name={weight_name}, adapter_name={adapter_name}\n")  
-                            elif weight_name and adapter_name==None and lora_weight_set == False:
-                                pipe.load_lora_weights(
-                                    lora_weight,
-                                    weight_name=weight_name,
-                                    token=constants.HF_API_TOKEN
-                                )
-                                lora_weight_set = True
-                                print(f"\npipe.load_lora_weights({lora_weight}, weight_name={weight_name}, adapter_name={adapter_name}\n")  
-                            elif lora_weight_set == False:
-                                pipe.load_lora_weights(
-                                    lora_weight,
-                                    token=constants.HF_API_TOKEN
-                                )  
-                                lora_weight_set = True
-                                print(f"\npipe.load_lora_weights({lora_weight}, weight_name={weight_name}, adapter_name={adapter_name}\n")  
-                        # Apply 'pipe' configurations if present
-                        if 'pipe' in config:
-                            pipe_config = config['pipe']
-                            for method_name, params in pipe_config.items():
-                                method = getattr(pipe, method_name, None)
-                                if method:
-                                    print(f"Applying pipe method: {method_name} with params: {params}")
-                                    method(**params)
-                                else:
-                                    print(f"Method {method_name} not found in pipe.")
-                        if 'condition_type' in config:
-                            condition_type = config['condition_type']
-                            if condition_type == "coloring":
-                                #pipe.enable_coloring()
-                                print("\nEnabled coloring.\n")
-                            elif condition_type == "deblurring":
-                                #pipe.enable_deblurring()
-                                print("\nEnabled deblurring.\n")
-                            elif condition_type == "fill":
-                                #pipe.enable_fill()
-                                print("\nEnabled fill.\n")
-                            elif condition_type == "depth":
-                                #pipe.enable_depth()
-                                print("\nEnabled depth.\n")
-                            elif condition_type == "canny":
-                                #pipe.enable_canny()
-                                print("\nEnabled canny.\n")
-                            elif condition_type == "subject":
-                                #pipe.enable_subject()
-                                print("\nEnabled subject.\n")
+    # alternative version that may be more efficient
+    # pipe.enable_sequential_cpu_offload()
+    if pipeline_name == "FluxPipeline":
+        pipe.enable_model_cpu_offload()
+        pipe.vae.enable_slicing()
+        #pipe.vae.enable_tiling()
+    else:
+        pipe.enable_model_cpu_offload()
+
+    # Access the tokenizer from the pipeline
+    tokenizer = pipe.tokenizer
+
+    # Check if add_prefix_space is set and convert to slow tokenizer if necessary
+    if getattr(tokenizer, 'add_prefix_space', False):
+        tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False, device_map = 'cpu')
+        # Update the pipeline's tokenizer
+        pipe.tokenizer = tokenizer
+        pipe.to(device)
+
+    flash_attention_enabled = torch.backends.cuda.flash_sdp_enabled()
+    if flash_attention_enabled == False:
+        #Enable xFormers memory-efficient attention (optional)
+        #pipe.enable_xformers_memory_efficient_attention()
+        print("\nEnabled xFormers memory-efficient attention.\n")
+    else:            
+        pipe.attn_implementation="flash_attention_2"
+        print("\nEnabled flash_attention_2.\n")
+
+    condition_type = "subject"
+    # Load LoRA weights
+    # note: does not yet handle multiple LoRA weights with different names, needs .set_adapters(["depth", "hyper-sd"], adapter_weights=[0.85, 0.125])        
+    if lora_weights:
+        for lora_weight in lora_weights:
+            lora_configs = constants.LORA_DETAILS.get(lora_weight, [])
+            lora_weight_set = False
+            if lora_configs:
+                for config in lora_configs:
+                    # Load LoRA weights with optional weight_name and adapter_name
+                    if 'weight_name' in config:
+                        weight_name = config.get("weight_name")
+                        adapter_name = config.get("adapter_name")
+                        lora_collection = config.get("lora_collection")
+                        if weight_name and adapter_name and lora_collection and lora_weight_set == False:
+                            pipe.load_lora_weights(
+                                lora_collection,
+                                weight_name=weight_name,
+                                adapter_name=adapter_name,
+                                token=constants.HF_API_TOKEN
+                            )
+                            lora_weight_set = True
+                            print(f"\npipe.load_lora_weights({lora_weight}, weight_name={weight_name}, adapter_name={adapter_name}, lora_collection={lora_collection}\n")  
+                        elif weight_name and adapter_name==None and lora_collection and lora_weight_set == False:
+                            pipe.load_lora_weights(
+                                lora_collection,
+                                weight_name=weight_name,
+                                token=constants.HF_API_TOKEN
+                            )
+                            lora_weight_set = True
+                            print(f"\npipe.load_lora_weights({lora_weight}, weight_name={weight_name}, adapter_name={adapter_name}, lora_collection={lora_collection}\n")  
+                        elif weight_name and adapter_name and lora_weight_set == False:
+                            pipe.load_lora_weights(
+                                lora_weight,
+                                weight_name=weight_name,
+                                adapter_name=adapter_name,
+                                token=constants.HF_API_TOKEN
+                            )
+                            lora_weight_set = True
+                            print(f"\npipe.load_lora_weights({lora_weight}, weight_name={weight_name}, adapter_name={adapter_name}\n")  
+                        elif weight_name and adapter_name==None and lora_weight_set == False:
+                            pipe.load_lora_weights(
+                                lora_weight,
+                                weight_name=weight_name,
+                                token=constants.HF_API_TOKEN
+                            )
+                            lora_weight_set = True
+                            print(f"\npipe.load_lora_weights({lora_weight}, weight_name={weight_name}, adapter_name={adapter_name}\n")  
+                        elif lora_weight_set == False:
+                            pipe.load_lora_weights(
+                                lora_weight,
+                                token=constants.HF_API_TOKEN
+                            )  
+                            lora_weight_set = True
+                            print(f"\npipe.load_lora_weights({lora_weight}, weight_name={weight_name}, adapter_name={adapter_name}\n")  
+                    # Apply 'pipe' configurations if present
+                    if 'pipe' in config:
+                        pipe_config = config['pipe']
+                        for method_name, params in pipe_config.items():
+                            method = getattr(pipe, method_name, None)
+                            if method:
+                                print(f"Applying pipe method: {method_name} with params: {params}")
+                                method(**params)
                             else:
-                                print(f"Condition type {condition_type} not implemented.")
-                else:
-                    pipe.load_lora_weights(lora_weight, use_auth_token=constants.HF_API_TOKEN)
-        # Set the random seed for reproducibility
-        generator = torch.Generator(device=device).manual_seed(seed)
-        conditions = []
-        if conditioned_image is not None:
-            conditioned_image = crop_and_resize_image(conditioned_image, image_width, image_height)
-            condition = Condition(condition_type, conditioned_image)
-            conditions.append(condition)
-            print(f"\nAdded conditioned image.\n {conditioned_image.size}")
-            # Prepare the parameters for image generation
-            additional_parameters ={
-                "strength": strength,
-                "image": conditioned_image,
-            }
-        else:
-            print("\nNo conditioned image provided.")
-            if neg_prompt!=None:
-                true_cfg_scale=1.1
-            additional_parameters ={
-                "negative_prompt": neg_prompt,
-                "true_cfg_scale": true_cfg_scale,
-            }
-        # handle long prompts by splitting them
-        if approximate_token_count(text) > 76:
-            prompt, prompt2 = split_prompt_precisely(text)
-            prompt_parameters = {
-                "prompt" : prompt,
-                "prompt_2": prompt2
-            }
-        else:
-            prompt_parameters = {
-                "prompt" :text
+                                print(f"Method {method_name} not found in pipe.")
+                    if 'condition_type' in config:
+                        condition_type = config['condition_type']
+                        if condition_type == "coloring":
+                            #pipe.enable_coloring()
+                            print("\nEnabled coloring.\n")
+                        elif condition_type == "deblurring":
+                            #pipe.enable_deblurring()
+                            print("\nEnabled deblurring.\n")
+                        elif condition_type == "fill":
+                            #pipe.enable_fill()
+                            print("\nEnabled fill.\n")
+                        elif condition_type == "depth":
+                            #pipe.enable_depth()
+                            print("\nEnabled depth.\n")
+                        elif condition_type == "canny":
+                            #pipe.enable_canny()
+                            print("\nEnabled canny.\n")
+                        elif condition_type == "subject":
+                            #pipe.enable_subject()
+                            print("\nEnabled subject.\n")
+                        else:
+                            print(f"Condition type {condition_type} not implemented.")
+            else:
+                pipe.load_lora_weights(lora_weight, use_auth_token=constants.HF_API_TOKEN)
+            gr.Info("lora_weights are loaded",duration=5)
+    # Set the random seed for reproducibility
+    generator = torch.Generator(device=device).manual_seed(seed)
+    conditions = []
+    if conditioned_image is not None:
+        conditioned_image = crop_and_resize_image(conditioned_image, image_width, image_height)
+        condition = Condition(condition_type, conditioned_image)
+        conditions.append(condition)
+        print(f"\nAdded conditioned image.\n {conditioned_image.size}")
+        # Prepare the parameters for image generation
+        additional_parameters ={
+            "strength": strength,
+            "image": conditioned_image,
         }
-        additional_parameters.update(prompt_parameters)
-        # Combine all parameters
-        generate_params = {
-            "height": image_height,
-            "width": image_width,
-            "guidance_scale": guidance_scale,
-            "num_inference_steps": num_inference_steps,
-            "generator": generator,        }
-        if additional_parameters:
-            generate_params.update(additional_parameters)
-        generate_params = {k: v for k, v in generate_params.items() if v is not None}
-        print(f"generate_params: {generate_params}")
-        # Generate the image
-        result = pipe(**generate_params) #generate_image(pipe,generate_params)
-        image = result.images[0]
-        # Clean up
-        del result
-        del conditions
-        del generator
-    # Delete the pipeline and clear cache
-    del pipe
-    torch.cuda.empty_cache()
-    torch.cuda.ipc_collect()
-    print(torch.cuda.memory_summary(device=None, abbreviated=False))
-    
-    return image
+    else:
+        print("\nNo conditioned image provided.")
+        if neg_prompt!=None:
+            true_cfg_scale=1.1
+        additional_parameters ={
+            "negative_prompt": neg_prompt,
+            "true_cfg_scale": true_cfg_scale,
+        }
+    # handle long prompts by splitting them
+    if approximate_token_count(text) > 76:
+        prompt, prompt2 = split_prompt_precisely(text)
+        prompt_parameters = {
+            "prompt" : prompt,
+            "prompt_2": prompt2
+        }
+    else:
+        prompt_parameters = {
+            "prompt" :text
+    }
+    additional_parameters.update(prompt_parameters)
+    # Combine all parameters
+    generate_params = {
+        "height": image_height,
+        "width": image_width,
+        "guidance_scale": guidance_scale,
+        "num_inference_steps": num_inference_steps,
+        "generator": generator,
+        }
+    if additional_parameters:
+        generate_params.update(additional_parameters)
+    generate_params = {k: v for k, v in generate_params.items() if v is not None}
+    print(f"generate_params: {generate_params}")
+    import pickle
+
+    try:
+        pickle.dumps(pipe)
+        print("pipe is picklable.\n")
+    except pickle.PicklingError:
+        print("pipe is not picklable\n.")
+
+    try:
+        pickle.dumps(conditions)
+        print("conditions is picklable.\n")
+    except pickle.PicklingError:
+        print("conditions is not picklable.\n")
+
+    try:
+        pickle.dumps(generator)
+        print("generator is picklable.\n")
+    except pickle.PicklingError:
+        print("generator is not picklable.\n")
+        
+    return pipe, conditions, generate_params
+
 
 def generate_ai_image_local (
     map_option,
@@ -565,7 +488,7 @@ def generate_ai_image_local (
         width = additional_parameters.pop('width', width)
         num_inference_steps = additional_parameters.pop('num_inference_steps', num_inference_steps)
         guidance_scale = additional_parameters.pop('guidance_scale', guidance_scale)
-        print("Generating image with the following parameters:")
+        print("Generating image with the following parameters:\n")
         print(f"Model: {model}")
         print(f"LoRA Weights: {lora_weights}")
         print(f"Prompt: {prompt}")
@@ -578,8 +501,8 @@ def generate_ai_image_local (
         print(f"Additional Parameters: {additional_parameters}")
         print(f"Conditioned Image: {conditioned_image}")
         print(f"Conditioned Image Strength: {strength}")
-        print(f"pipeline: {pipeline_name}")
-        image = generate_image_lowmem(
+        print(f"pipeline: {pipeline_name}\n")
+        pipe, conditions, generate_params = generate_image_lowmem(
             text=prompt,
             model_name=model,
             neg_prompt=negative_prompt,
@@ -594,6 +517,7 @@ def generate_ai_image_local (
             strength=strength,
             additional_parameters=additional_parameters
         )
+        image = generate_image(pipe, conditions, **generate_params)
         with NamedTemporaryFile(delete=False, suffix=".png") as tmp:
             image.save(tmp.name, format="PNG")
             constants.temp_files.append(tmp.name)
@@ -684,7 +608,8 @@ def on_prerendered_gallery_selection(event_data: gr.SelectData):
     global current_prerendered_image
     selected_index = event_data.index
     selected_image = constants.pre_rendered_maps_paths[selected_index]
-    print(f"Gallery Image Selected: {selected_image}\n")
+    print(f"Template Image Selected: {selected_image} ({event_data.index})\n")
+    gr.Info(f"Template Image Selected: {selected_image} ({event_data.index})",duration=5)
     current_prerendered_image.value = selected_image
     return current_prerendered_image
 
@@ -1021,11 +946,11 @@ with gr.Blocks(css_paths="style_20250128.css", title=title, theme='Surn/beeuty',
                 )
 
     with gr.Row():
-        with gr.Accordion("Generate AI Image (click here)", open = False):
+        with gr.Accordion("Generate AI Image (click here for options)", open = False):
             with gr.Row():
                 with gr.Column():
                     model_options = gr.Dropdown(
-                        label="Model Options",
+                        label="Choose an AI Model*",
                         choices=constants.MODELS + constants.LORA_WEIGHTS + ["Manual Entry"],
                         value="Cossale/Frames2-Flex.1",
                         elem_classes="solid"

utils/constants.py

@@ -54,6 +54,7 @@ def load_env_vars(env_path):
 # os.environ['TMPDIR'] = r'e:\\TMP'
 # os.environ['XDG_CACHE_HOME'] = r'E:\\cache'
 
+USE_FLASH_ATTENTION = os.getenv("USE_FLASH_ATTENTION", "0") == "1"
 HF_API_TOKEN = os.getenv("HF_TOKEN")
 if not HF_API_TOKEN:
     raise ValueError("HF_TOKEN is not set. Please check your .env file.")