Fall Back from update

app.py

@@ -8,7 +8,6 @@ from typing import Optional, Union, List, Tuple
 
 from PIL import Image, ImageFilter
 import cv2
-
 import utils.constants as constants
 
 from haishoku.haishoku import Haishoku
@@ -92,7 +91,6 @@ from utils.version_info import (
     #release_torch_resources,
     #get_torch_info
 )
-from src.condition import Condition
 import spaces
 
 input_image_palette = []
@@ -201,24 +199,11 @@ condition_dict = {
     "fill": 9,
 }
 
-@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=90)
-@torch.no_grad()
+@spaces.GPU(duration=200, progress=gr.Progress(track_tqdm=True))
 def generate_image_lowmem(
     text,
     neg_prompt=None,
@@ -244,205 +229,195 @@ def generate_image_lowmem(
                         f"Available options: {list(PIPELINE_CLASSES.keys())}")
 
     #initialize_cuda()
-    device = "cuda" if torch.cuda.is_available() else "cpu"    
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    from src.condition import Condition
 
     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():
-    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
+    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
         
-    # 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)
+        # 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")
                             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")
-                        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,
-        }
-    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,
+                                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
         }
-    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
-
+        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
 
 def generate_ai_image_local (
     map_option,
@@ -501,8 +476,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}\n")
-        pipe, conditions, generate_params = generate_image_lowmem(
+        print(f"pipeline: {pipeline_name}")
+        image = generate_image_lowmem(
             text=prompt,
             model_name=model,
             neg_prompt=negative_prompt,
@@ -517,7 +492,6 @@ 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)