Reverse Changes

app.py

@@ -1,5 +1,7 @@
 import os
 # Import constants
+import numpy as np
+import torch
 import utils.constants as constants
 import gradio as gr
 from PIL import Image
@@ -10,6 +12,9 @@ from tempfile import NamedTemporaryFile
 import atexit
 import random
 import logging
+import accelerate 
+from transformers import AutoTokenizer
+import gc
 
 IS_SHARED_SPACE = constants.IS_SHARED_SPACE
 
@@ -57,13 +62,25 @@ from utils.excluded_colors import (
 # )
 from utils.version_info import (
     versions_html,
-    initialize_cuda,
-    release_torch_resources,
+    #initialize_cuda,
+    #release_torch_resources,
     get_torch_info
 )
 from utils.lora_details import (
-    upd_prompt_notes
+    upd_prompt_notes,
+    split_prompt_precisely,
+    approximate_token_count,
+    get_trigger_words
 )
+from diffusers import FluxPipeline,FluxImg2ImgPipeline,FluxControlPipeline
+
+PIPELINE_CLASSES = {
+    "FluxPipeline": FluxPipeline,
+    "FluxImg2ImgPipeline": FluxImg2ImgPipeline,
+    "FluxControlPipeline": FluxControlPipeline
+}
+
+import spaces
 
 input_image_palette = []
 current_prerendered_image = gr.State("./images/images/Beeuty-1.png")
@@ -162,324 +179,313 @@ def get_model_and_lora(model_textbox):
         default_model = model_textbox
         return default_model, []
 
-
-def generate_input_image_click(map_option, prompt_textbox_value, negative_prompt_textbox_value, model_textbox_value, randomize_seed=True, seed=None, use_conditioned_image=False, strength=0.5, image_format="16:9", scale_factor=(8/3), progress=gr.Progress(track_tqdm=True)):
-    import spaces
-    from diffusers import FluxPipeline,FluxImg2ImgPipeline,FluxControlPipeline
-    import accelerate 
-    from transformers import AutoTokenizer
-    from utils.lora_details import get_trigger_words, approximate_token_count, split_prompt_precisely
-    import gc
-    PIPELINE_CLASSES = {
-        "FluxPipeline": FluxPipeline,
-        "FluxImg2ImgPipeline": FluxImg2ImgPipeline,
-        "FluxControlPipeline": FluxControlPipeline
-    }
-    if randomize_seed:
-        seed = random.randint(0, constants.MAX_SEED)
-
-    @spaces.GPU(progress=gr.Progress(track_tqdm=True))
-    def generate_image_lowmem(
-        text,
-        neg_prompt=None,
-        model_name="black-forest-labs/FLUX.1-dev",
-        lora_weights=None,
-        conditioned_image=None,
-        image_width=1368,
-        image_height=848,
-        guidance_scale=3.5,
-        num_inference_steps=30,
-        seed=0,
-        true_cfg_scale=1.0,
-        pipeline_name="FluxPipeline",
-        strength=0.75,
-        additional_parameters=None,
-        progress=gr.Progress(track_tqdm=True)
-    ):  
-        from torch import cuda, bfloat16, float32, Generator, no_grad, backends
-        # Retrieve the pipeline class from the mapping
-        pipeline_class = PIPELINE_CLASSES.get(pipeline_name)
-        if not pipeline_class:
-            raise ValueError(f"Unsupported pipeline type '{pipeline_name}'. "
-                            f"Available options: {list(PIPELINE_CLASSES.keys())}")
-
-        #initialize_cuda()
-        device = "cuda" if 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 no_grad():
-            # Initialize the pipeline inside the context manager
-            pipe = pipeline_class.from_pretrained(
-                model_name,
-                torch_dtype=bfloat16 if device == "cuda" else float32
-            ).to(device)
-            # Optionally, don't use CPU offload if not necessary
+@spaces.GPU(progress=gr.Progress(track_tqdm=True))
+def generate_image_lowmem(
+    text,
+    neg_prompt=None,
+    model_name="black-forest-labs/FLUX.1-dev",
+    lora_weights=None,
+    conditioned_image=None,
+    image_width=1368,
+    image_height=848,
+    guidance_scale=3.5,
+    num_inference_steps=30,
+    seed=0,
+    true_cfg_scale=1.0,
+    pipeline_name="FluxPipeline",
+    strength=0.75,
+    additional_parameters=None,
+    progress=gr.Progress(track_tqdm=True)
+):  
+    #from torch import cuda, bfloat16, float32, Generator, no_grad, backends
+    # Retrieve the pipeline class from the mapping
+    pipeline_class = PIPELINE_CLASSES.get(pipeline_name)
+    if not pipeline_class:
+        raise ValueError(f"Unsupported pipeline type '{pipeline_name}'. "
+                        f"Available options: {list(PIPELINE_CLASSES.keys())}")
+
+    #initialize_cuda()
+    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():
+        # 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 = 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 = 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)
+        # 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
             }
-            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)
-            image = result.images[0]
-            # Clean up
-            del result
-            del conditions
-            del generator
-        # Delete the pipeline and clear cache
-        del pipe
-        cuda.empty_cache()
-        cuda.ipc_collect()
-        print(cuda.memory_summary(device=None, abbreviated=False))
+        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}")
+        # Generate the image
+        result = 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
-
-    #@spaces.GPU(progress=gr.Progress(track_tqdm=True))
-    def generate_ai_image_local (
-        map_option,
-        prompt_textbox_value,
-        neg_prompt_textbox_value,
-        model="black-forest-labs/FLUX.1-dev",
-        lora_weights=None,
-        conditioned_image=None,
-        height=512,
-        width=912,
-        num_inference_steps=30,
-        guidance_scale=3.5,
-        seed=777,
-        pipeline_name="FluxPipeline",
-        strength=0.75,
-        progress=gr.Progress(track_tqdm=True)
-    ):        
-        print(f"Generating image with lowmem")
-        try:
-            if map_option != "Prompt":
-                prompt = constants.PROMPTS[map_option]
-                negative_prompt = constants.NEGATIVE_PROMPTS.get(map_option, "")
-            else:
-                prompt = prompt_textbox_value
-                negative_prompt = neg_prompt_textbox_value or ""
-            #full_prompt = f"{prompt} {negative_prompt}"
-            additional_parameters = {}
-            if lora_weights:
-                for lora_weight in lora_weights:
-                    lora_configs = constants.LORA_DETAILS.get(lora_weight, [])
-                    for config in lora_configs:
-                        if 'parameters' in config:
-                            additional_parameters.update(config['parameters'])
-                        elif 'trigger_words' in config:
-                            trigger_words = get_trigger_words(lora_weight)
-                            prompt = f"{trigger_words} {prompt}"
-            for key, value in additional_parameters.items():
-                if key in ['height', 'width', 'num_inference_steps', 'max_sequence_length']:
-                    additional_parameters[key] = int(value)
-                elif key in ['guidance_scale','true_cfg_scale']:
-                    additional_parameters[key] = float(value)
-            height = additional_parameters.pop('height', height)
-            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(f"Model: {model}")
-            print(f"LoRA Weights: {lora_weights}")
-            print(f"Prompt: {prompt}")
-            print(f"Neg Prompt: {negative_prompt}")
-            print(f"Height: {height}")
-            print(f"Width: {width}")
-            print(f"Number of Inference Steps: {num_inference_steps}")
-            print(f"Guidance Scale: {guidance_scale}")
-            print(f"Seed: {seed}")
-            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(
-                text=prompt,
-                model_name=model,
-                neg_prompt=negative_prompt,
-                lora_weights=lora_weights,
-                conditioned_image=conditioned_image,
-                image_width=width,
-                image_height=height,
-                guidance_scale=guidance_scale,
-                num_inference_steps=num_inference_steps,
-                seed=seed,
-                pipeline_name=pipeline_name,
-                strength=strength,
-                additional_parameters=additional_parameters
-            )
-            with NamedTemporaryFile(delete=False, suffix=".png") as tmp:
-                image.save(tmp.name, format="PNG")
-                constants.temp_files.append(tmp.name)
-                print(f"Image saved to {tmp.name}")
-                release_torch_resources()
-                gc.collect()
-                return tmp.name
-        except Exception as e:
-            print(f"Error generating AI image: {e}")
-            release_torch_resources()
+    return image
+
+def generate_ai_image_local (
+    map_option,
+    prompt_textbox_value,
+    neg_prompt_textbox_value,
+    model="black-forest-labs/FLUX.1-dev",
+    lora_weights=None,
+    conditioned_image=None,
+    height=512,
+    width=912,
+    num_inference_steps=30,
+    guidance_scale=3.5,
+    seed=777,
+    pipeline_name="FluxPipeline",
+    strength=0.75,
+    progress=gr.Progress(track_tqdm=True)
+):        
+    print(f"Generating image with lowmem")
+    try:
+        if map_option != "Prompt":
+            prompt = constants.PROMPTS[map_option]
+            negative_prompt = constants.NEGATIVE_PROMPTS.get(map_option, "")
+        else:
+            prompt = prompt_textbox_value
+            negative_prompt = neg_prompt_textbox_value or ""
+        #full_prompt = f"{prompt} {negative_prompt}"
+        additional_parameters = {}
+        if lora_weights:
+            for lora_weight in lora_weights:
+                lora_configs = constants.LORA_DETAILS.get(lora_weight, [])
+                for config in lora_configs:
+                    if 'parameters' in config:
+                        additional_parameters.update(config['parameters'])
+                    elif 'trigger_words' in config:
+                        trigger_words = get_trigger_words(lora_weight)
+                        prompt = f"{trigger_words} {prompt}"
+        for key, value in additional_parameters.items():
+            if key in ['height', 'width', 'num_inference_steps', 'max_sequence_length']:
+                additional_parameters[key] = int(value)
+            elif key in ['guidance_scale','true_cfg_scale']:
+                additional_parameters[key] = float(value)
+        height = additional_parameters.pop('height', height)
+        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(f"Model: {model}")
+        print(f"LoRA Weights: {lora_weights}")
+        print(f"Prompt: {prompt}")
+        print(f"Neg Prompt: {negative_prompt}")
+        print(f"Height: {height}")
+        print(f"Width: {width}")
+        print(f"Number of Inference Steps: {num_inference_steps}")
+        print(f"Guidance Scale: {guidance_scale}")
+        print(f"Seed: {seed}")
+        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(
+            text=prompt,
+            model_name=model,
+            neg_prompt=negative_prompt,
+            lora_weights=lora_weights,
+            conditioned_image=conditioned_image,
+            image_width=width,
+            image_height=height,
+            guidance_scale=guidance_scale,
+            num_inference_steps=num_inference_steps,
+            seed=seed,
+            pipeline_name=pipeline_name,
+            strength=strength,
+            additional_parameters=additional_parameters
+        )
+        with NamedTemporaryFile(delete=False, suffix=".png") as tmp:
+            image.save(tmp.name, format="PNG")
+            constants.temp_files.append(tmp.name)
+            print(f"Image saved to {tmp.name}")
+            #release_torch_resources()
             gc.collect()
-            return None
+            return tmp.name
+    except Exception as e:
+        print(f"Error generating AI image: {e}")
+        #release_torch_resources()
+        gc.collect()
+        return None
+
+@spaces.GPU(duration=140,progress=gr.Progress(track_tqdm=True))
+def generate_input_image_click(map_option, prompt_textbox_value, negative_prompt_textbox_value, model_textbox_value, randomize_seed=True, seed=None, use_conditioned_image=False, strength=0.5, image_format="16:9", scale_factor=(8/3), progress=gr.Progress(track_tqdm=True)):
+    if randomize_seed:
+        seed = random.randint(0, constants.MAX_SEED)
+
     # Get the model and LoRA weights
     model, lora_weights = get_model_and_lora(model_textbox_value)
     global current_prerendered_image
@@ -565,7 +571,12 @@ def add_border(image, mask_width, mask_height, blank_color):
     margin_color = detect_color_format(blank_color)
     print(f"Adding border to image with width: {mask_width}, height: {mask_height}, color: {margin_color}")
     return shrink_and_paste_on_blank(bordered_image_output, mask_width, mask_height, margin_color)
-def main(debug):
+
+@spaces.GPU()
+def getVersions():
+    return versions_html()
+generate_input_image_click.zerogpu = True
+def main(debug=False):
     title = "HexaGrid Creator"
     #description = "Customizable Hexagon Grid Image Generator"
     examples = [["assets//examples//hex_map_p1.png", 32, 1, 0, 0, 0, 0, 0, "#ede9ac44","#12165380", True]]
@@ -791,7 +802,7 @@ def main(debug):
                                 # Gallery from PRE_RENDERED_IMAGES GOES HERE
                                 prerendered_image_gallery = gr.Gallery(label="Image Gallery", show_label=True, value=build_prerendered_images(constants.pre_rendered_maps_paths), elem_id="gallery", elem_classes="solid", type="filepath", columns=[3], rows=[3], preview=False ,object_fit="contain", height="auto", format="png",allow_preview=False)
                              with gr.Row():
-                                 image_guidance_stength = gr.Slider(label="Image Guidance Strength", minimum=0, maximum=1.0, value=0.25, step=0.01, interactive=True)
+                                 image_guidance_stength = gr.Slider(label="Image Guidance Strength (prompt <-> image)", minimum=0, maximum=1.0, value=0.5, step=0.01, interactive=True)
                                  with gr.Column():
                                      replace_input_image_button = gr.Button(
                                          "Replace Input Image",
@@ -888,35 +899,35 @@ def main(debug):
                 ],
                 inputs=[input_image, filter_color, fill_hex, start_x, start_y, end_x, end_y, x_spacing, y_spacing, hex_size, rotation, border_size, border_color, border_opacity],
                 elem_id="examples")
-        #with gr.Row():
-            #gr.HTML(value=versions_html(), visible=True, elem_id="versions")
         with gr.Row():
-            reinstall_torch = gr.Button("Reinstall Torch", elem_classes="solid small", variant="secondary")
-            reinstall_cuda_toolkit = gr.Button("Install CUDA Toolkit", elem_classes="solid small", variant="secondary")
-            reinitialize_cuda = gr.Button("Reinitialize CUDA", elem_classes="solid small", variant="secondary")
-            torch_release = gr.Button("Release Torch Resources", elem_classes="solid small", variant="secondary")
-
-        reinitialize_cuda.click(
-            fn=initialize_cuda,
-            inputs=[],
-            outputs=[]
-        )
-        torch_release.click(
-            fn=release_torch_resources,
-            inputs=[],
-            outputs=[]
-        )
-        reinstall_torch.click(
-            fn=install_torch,
-            inputs=[],
-            outputs=[]
-        )
-
-        reinstall_cuda_toolkit.click(
-            fn=install_cuda_toolkit,
-            inputs=[],
-            outputs=[]
-        )
+            gr.HTML(value=getVersions(), visible=True, elem_id="versions")
+        # with gr.Row():
+        #     reinstall_torch = gr.Button("Reinstall Torch", elem_classes="solid small", variant="secondary")
+        #     reinstall_cuda_toolkit = gr.Button("Install CUDA Toolkit", elem_classes="solid small", variant="secondary")
+        #     reinitialize_cuda = gr.Button("Reinitialize CUDA", elem_classes="solid small", variant="secondary")
+        #     torch_release = gr.Button("Release Torch Resources", elem_classes="solid small", variant="secondary")
+
+        # reinitialize_cuda.click(
+        #     fn=initialize_cuda,
+        #     inputs=[],
+        #     outputs=[]
+        # )
+        # torch_release.click(
+        #     fn=release_torch_resources,
+        #     inputs=[],
+        #     outputs=[]
+        # )
+        # reinstall_torch.click(
+        #     fn=install_torch,
+        #     inputs=[],
+        #     outputs=[]
+        # )
+
+        # reinstall_cuda_toolkit.click(
+        #     fn=install_cuda_toolkit,
+        #     inputs=[],
+        #     outputs=[]
+        # )
 
         color_display.select(on_color_display_select,inputs=[color_display], outputs=[selected_row])
         color_display.input(on_input,inputs=[color_display], outputs=[color_display, gr.State(excluded_color_list)])
@@ -1004,8 +1015,8 @@ def main(debug):
             outputs=[bordered_image_output],
             scroll_to_output=True
         )
-    (())
-    beeuty.queue(default_concurrency_limit=1,max_size=12,api_open=False)
+
+    beeuty.queue(default_concurrency_limit=2,max_size=12,api_open=False)
     beeuty.launch(allowed_paths=["assets","/","./assets","images","./images", "./images/prerendered"], favicon_path="./assets/favicon.ico", max_file_size="10mb")
 
 
@@ -1014,12 +1025,13 @@ if __name__ == "__main__":
         format="[%(levelname)s] %(asctime)s %(message)s", level=logging.INFO
     )
     logging.info("Environment Variables: %s" % os.environ)
-    if _get_output(["nvcc", "--version"]) is None:
-        logging.info("Installing CUDA toolkit...")
-        install_cuda_toolkit()
-    else:
-        logging.info("Detected CUDA: %s" % _get_output(["nvcc", "--version"]))
-
-    logging.info("Installing CUDA extensions...")
-    setup_runtime_env()
-    main(os.getenv("DEBUG") == "1")
+    # if _get_output(["nvcc", "--version"]) is None:
+    #     logging.info("Installing CUDA toolkit...")
+    #     install_cuda_toolkit()
+    # else:
+    #     logging.info("Detected CUDA: %s" % _get_output(["nvcc", "--version"]))
+
+    # logging.info("Installing CUDA extensions...")
+    # setup_runtime_env()
+    #main(os.getenv("DEBUG") == "1")
+    main()

utils/ai_generator_diffusers_flux.py

@@ -5,7 +5,7 @@ import utils.constants as constants
 import gradio as gr
 from torch import __version__ as torch_version_, version, cuda, bfloat16, float32, Generator, no_grad, backends
 from diffusers import FluxPipeline,FluxImg2ImgPipeline,FluxControlPipeline
-#import accelerate 
+import accelerate 
 from transformers import AutoTokenizer
 import safetensors
 #import xformers
@@ -19,9 +19,9 @@ from utils.image_utils import (
 )
 from utils.version_info import (
     get_torch_info,
-    get_diffusers_version,
-    get_transformers_version,
-    get_xformers_version,
+    # get_diffusers_version,
+    # get_transformers_version,
+    # get_xformers_version,
     initialize_cuda,
     release_torch_resources
 )

utils/constants.py

@@ -40,6 +40,7 @@ if not HF_API_TOKEN:
 
 default_lut_example_img = "./LUT/daisy.jpg"
 MAX_SEED = np.iinfo(np.int32).max
+TARGET_SIZE = (2688,1536)
 
 PROMPTS = {
     "BorderBlack": "Top-down view of a hexagon-based alien map with black borders. Features rivers, mountains, volcanoes, and snow at top and bottom. Colors: light blue, green, tan, brown. No reflections or shadows. Partial hexes on edges are black.",

utils/version_info.py

@@ -106,7 +106,7 @@ def versions_html():
          • 
         transformers: {get_transformers_version()}
          • 
-        xformers: {get_xformers_version()}
+        safetensors: {get_safetensors_version()}
          • 
         gradio: {gr.__version__}
          • 

web-ui.bat

@@ -1,2 +1,5 @@
-py -m app
+set NVIDIA_VISIBLE_DEVICES=0
+set CUDA_VISIBLE_DEVICES=0
+set CUDA_MODULE_LOADING=LAZY
+python -m app
 pause