Change Torch references

README.md

@@ -5,7 +5,7 @@ colorFrom: yellow
 colorTo: purple
 sdk: gradio
 python_version: 3.10.13
-sdk_version: 5.15.0
+sdk_version: 5.16.0
 app_file: app.py
 pinned: false
 short_description: Transform Your Images into Mesmerizing Hexagon Grids

app.py

@@ -1,14 +1,15 @@
-import spaces
+import os
+# Import constants
+import utils.constants as constants
 import gradio as gr
 from PIL import Image
 from haishoku.haishoku import Haishoku
-#import os
+
 from tempfile import NamedTemporaryFile
 #from pathlib import Path
 import atexit
 import random
-# Import constants
-import utils.constants as constants
+import logging
 
 IS_SHARED_SPACE = constants.IS_SHARED_SPACE
 
@@ -20,7 +21,7 @@ from utils.color_utils import (
     detect_color_format,
     update_color_opacity,
 )
-from utils.misc import (get_filename, pause, convert_ratio_to_dimensions)
+from utils.misc import (get_filename, pause, convert_ratio_to_dimensions, install_cuda_toolkit,install_torch, _get_output, setup_runtime_env)
 from utils.depth_estimation import estimate_depth, create_3d_model, generate_depth_and_3d, generate_depth_button_click
 
 from utils.image_utils import (
@@ -33,7 +34,8 @@ from utils.image_utils import (
     show_lut,
     apply_lut_to_image_path,
     multiply_and_blend_images,
-    alpha_composite_with_control
+    alpha_composite_with_control,
+    crop_and_resize_image
 )
 
 from utils.hex_grid import (
@@ -50,11 +52,13 @@ from utils.excluded_colors import (
     on_color_display_select
 )
 
-from utils.ai_generator import (
-    generate_ai_image,
-)
+# from utils.ai_generator import (
+#     generate_ai_image,
+# )
 from utils.version_info import (
     versions_html,
+    initialize_cuda,
+    release_torch_resources,
     get_torch_info
 )
 from utils.lora_details import (
@@ -138,8 +142,7 @@ def hex_create(hex_size, border_size, input_image_path, start_x, start_y, end_x,
         add_hex_text_option,
         custom_text_list,
         custom_text_color_list
-    )
-
+    )    
     return grid_image
 
 def get_model_and_lora(model_textbox):
@@ -159,7 +162,324 @@ 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, seed=None, use_conditioned_image=False, strength=0.5, image_format="16:9", scale_factor=3, 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)):
+    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
+        
+            # 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")
+                                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
+            }
+            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))
+    
+        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()
+            gc.collect()
+            return None
     # Get the model and LoRA weights
     model, lora_weights = get_model_and_lora(model_textbox_value)
     global current_prerendered_image
@@ -168,7 +488,7 @@ def generate_input_image_click(map_option, prompt_textbox_value, negative_prompt
     if use_conditioned_image:
         print(f"Conditioned path: {current_prerendered_image.value}.. converting to RGB\n")
         # ensure the conditioned image is an image and not a string, cannot use RGBA
-        if isinstance(current_prerendered_image.value, str):            
+        if isinstance(current_prerendered_image.value, str):
             conditioned_image = open_image(current_prerendered_image.value).convert("RGB")
             print(f"Conditioned Image: {conditioned_image.size}.. converted to RGB\n")
 
@@ -176,20 +496,23 @@ def generate_input_image_click(map_option, prompt_textbox_value, negative_prompt
     width_ratio, height_ratio = map(int, image_format.split(":"))
     aspect_ratio = width_ratio / height_ratio
     
-    width, height = convert_ratio_to_dimensions(aspect_ratio, 512)
-    
+    width, height = convert_ratio_to_dimensions(aspect_ratio, 576)
+    pipeline = "FluxPipeline"
+    if conditioned_image is not None:
+        pipeline = "FluxImg2ImgPipeline"
     # Generate the AI image and get the image path
-    image_path = generate_ai_image(
+    image_path = generate_ai_image_local(
         map_option,
         prompt_textbox_value,
         negative_prompt_textbox_value,
         model,
         lora_weights,
         conditioned_image,
-        stength=strength,
+        strength=strength,
         height=height,
         width=width,
-        seed=seed
+        seed=seed,
+        pipeline_name=pipeline,
     )
        
     # Open the generated image
@@ -207,7 +530,7 @@ def generate_input_image_click(map_option, prompt_textbox_value, negative_prompt
         upscaled_image.save(tmp_upscaled.name, format="PNG")
         constants.temp_files.append(tmp_upscaled.name)
         print(f"Upscaled image saved to {tmp_upscaled.name}")
-       
+    
     # Return the path of the upscaled image
     return tmp_upscaled.name
 
@@ -236,423 +559,467 @@ def combine_images_with_lerp(input_image, output_image, alpha):
     print(f"Combining images with alpha: {alpha}")
     return lerp_imagemath(in_image, out_image, alpha)
 
-def add_border(image, mask_width, mask_height, blank_color):
+def add_border(image, mask_width, mask_height, blank_color):    
+    #install_torch()
     bordered_image_output = Image.open(image).convert("RGBA")
     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):
+    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]]
 
-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]]
+    gr.set_static_paths(paths=["images/","images/images","images/prerendered","LUT/","fonts/"])
 
-gr.set_static_paths(paths=["images/","images/images","images/prerendered","LUT/","fonts/"])
-# Gradio Blocks Interface
-with gr.Blocks(css_paths="style_20250128.css", title="HexaGrid Creator", theme='Surn/beeuty') as beeuty:
-    with gr.Row():
-        gr.Markdown("""
-        # HexaGrid Creator
-        ## Transform Your Images into Mesmerizing Hexagon Grid Masterpieces! ⬢""", elem_classes="intro")
-    with gr.Row():
-        with gr.Accordion("Welcome to HexaGrid Creator, the ultimate tool for transforming your images into stunning hexagon grid artworks. Whether you're a tabletop game enthusiast, a digital artist, or someone who loves unique patterns, HexaGrid Creator has something for you.", open=False, elem_classes="intro"):
-            gr.Markdown ("""
+    # Gradio Blocks Interface
+    with gr.Blocks(css_paths="style_20250128.css", title=title, theme='Surn/beeuty') as beeuty:
+        with gr.Row():
+            gr.Markdown("""
+            # HexaGrid Creator
+            ## Transform Your Images into Mesmerizing Hexagon Grid Masterpieces! ⬢""", elem_classes="intro")
+        with gr.Row():
+            with gr.Accordion("Welcome to HexaGrid Creator, the ultimate tool for transforming your images into stunning hexagon grid artworks. Whether you're a tabletop game enthusiast, a digital artist, or someone who loves unique patterns, HexaGrid Creator has something for you.", open=False, elem_classes="intro"):
+                gr.Markdown ("""
 
-            ## Drop an image into the Input Image and get started!
+                ## Drop an image into the Input Image and get started!
 
         
 
-            ## What is HexaGrid Creator?
-            HexaGrid Creator is a web-based application that allows you to apply a hexagon grid overlay to any image. You can customize the size, color, and opacity of the hexagons, as well as the background and border colors. The result is a visually striking image that looks like it was made from hexagonal tiles!
-
-            ### What Can You Do?
-            - **Generate Hexagon Grids:** Create beautiful hexagon grid overlays on any image with fully customizable parameters.
-            - **AI-Powered Image Generation:** Use advanced AI models to generate images based on your prompts and apply hexagon grids to them.
-            - **Color Exclusion:** Select and exclude specific colors from your hexagon grid for a cleaner and more refined look.
-            - **Interactive Customization:** Adjust hexagon size, border size, rotation, background color, and more in real-time.
-            - **Depth and 3D Model Generation:** Generate depth maps and 3D models from your images for enhanced visualization.
-            - **Image Filter [Look-Up Table (LUT)] Application:** Apply filters (LUTs) to your images for color grading and enhancement.
-            - **Pre-rendered Maps:** Access a library of pre-rendered hexagon maps for quick and easy customization.
-            - **Add Margins:** Add customizable margins around your images for a polished finish.
-
-            ### Why You'll Love It
-            - **Fun and Easy to Use:** With an intuitive interface and real-time previews, creating hexagon grids has never been this fun!
-            - **Endless Creativity:** Unleash your creativity with endless customization options and see your images transform in unique ways.
-            - **Hexagon-Inspired Theme:** Enjoy a delightful yellow and purple theme inspired by hexagons! ⬢
-            - **Advanced AI Models:** Leverage advanced AI models and LoRA weights for high-quality image generation and customization.
-
-            ### Get Started
-            1. **Upload or Generate an Image:** Start by uploading your own image or generate one using our AI-powered tool.
-            2. **Customize Your Grid:** Play around with the settings to create the perfect hexagon grid overlay.
-            3. **Download and Share:** Once you're happy with your creation, download it and share it with the world!
-
-            ### Advanced Features
-            - **Generative AI Integration:** Utilize models like `black-forest-labs/FLUX.1-dev` and various LoRA weights for generating unique images.
-            - **Pre-rendered Maps:** Access a library of pre-rendered hexagon maps for quick and easy customization.
-            - **Image Filter [Look-Up Table (LUT)] Application:** Apply filters (LUTs) to your images for color grading and enhancement.
-            - **Depth and 3D Model Generation:** Create depth maps and 3D models from your images for enhanced visualization.
-            - **Add Margins:** Customize margins around your images for a polished finish.
-
-            Join the hive and start creating with HexaGrid Creator today!
-        
-            """, elem_classes="intro")
-    with gr.Row():
-        from utils.image_utils import convert_to_rgba_png
-
-        # Existing code
-        with gr.Column(scale=2):
-            input_image = gr.Image(
-                label="Input Image",
-                type="filepath",
-                interactive=True,
-                elem_classes="centered solid imgcontainer",
-                key="imgInput",
-                image_mode=None,
-                format="PNG", 
-                show_download_button=True,
-            )
+                ## What is HexaGrid Creator?
+                HexaGrid Creator is a web-based application that allows you to apply a hexagon grid overlay to any image. You can customize the size, color, and opacity of the hexagons, as well as the background and border colors. The result is a visually striking image that looks like it was made from hexagonal tiles!
 
-        # New code to convert input image to RGBA PNG
-        def on_input_image_change(image_path):
-            if image_path is None:
-                gr.Warning("Please upload an Input Image to get started.")
-                return None
-            img, img_path = convert_to_rgba_png(image_path)
-            return img_path
-
-        input_image.change(
-            fn=on_input_image_change,
-            inputs=[input_image],
-            outputs=[input_image], scroll_to_output=True,
-        )
-        with gr.Column():
-            with gr.Accordion("Hex Coloring and Exclusion", open = False):
-                with gr.Row():
-                    with gr.Column():
-                        color_picker = gr.ColorPicker(label="Pick a color to exclude",value="#505050")
-                    with gr.Column():
-                        filter_color = gr.Checkbox(label="Filter Excluded Colors from Sampling", value=False,)
-                exclude_color_button = gr.Button("Exclude Color", elem_id="exlude_color_button", elem_classes="solid")
-                color_display = gr.DataFrame(label="List of Excluded RGBA Colors", headers=["R", "G", "B", "A"], elem_id="excluded_colors", type="array", value=build_dataframe(excluded_color_list), interactive=True, elem_classes="solid centered")
-                selected_row = gr.Number(0, label="Selected Row", visible=False)
-                delete_button = gr.Button("Delete Row", elem_id="delete_exclusion_button", elem_classes="solid")
-                fill_hex = gr.Checkbox(label="Fill Hex with color from Image", value=True)
-            with gr.Accordion("Image Filters", open = False):
-                with gr.Row():
-                    with gr.Column():
-                        composite_color = gr.ColorPicker(label="Color", value="#ede9ac44")
-                    with gr.Column():
-                        composite_opacity = gr.Slider(label="Opacity %", minimum=0, maximum=100, value=50, interactive=True)
-                with gr.Row():
-                    composite_button = gr.Button("Composite", elem_classes="solid")
-                with gr.Row():
-                    with gr.Column():
-                        lut_filename = gr.Textbox(
-                            value="", 
-                            label="Look Up Table (LUT) File Name",
-                            elem_id="lutFileName")
-                    with gr.Column():
-                        lut_file = gr.File(
-                            value=None,
-                            file_count="single",
-                            file_types=[".cube"],
-                            type="filepath",
-                            label="LUT cube File")
-                with gr.Row():
-                    lut_example_image = gr.Image(type="pil", label="Filter (LUT) Example Image", value=constants.default_lut_example_img)
-                with gr.Row():
-                    with gr.Column():
-                        gr.Markdown("""
-                        ### Included Filters (LUTs)
-                        There are several included Filters:
+                ### What Can You Do?
+                - **Generate Hexagon Grids:** Create beautiful hexagon grid overlays on any image with fully customizable parameters.
+                - **AI-Powered Image Generation:** Use advanced AI models to generate images based on your prompts and apply hexagon grids to them.
+                - **Color Exclusion:** Select and exclude specific colors from your hexagon grid for a cleaner and more refined look.
+                - **Interactive Customization:** Adjust hexagon size, border size, rotation, background color, and more in real-time.
+                - **Depth and 3D Model Generation:** Generate depth maps and 3D models from your images for enhanced visualization.
+                - **Image Filter [Look-Up Table (LUT)] Application:** Apply filters (LUTs) to your images for color grading and enhancement.
+                - **Pre-rendered Maps:** Access a library of pre-rendered hexagon maps for quick and easy customization.
+                - **Add Margins:** Add customizable margins around your images for a polished finish.
 
-                        Try them on the example image before applying to your Input Image.
-                        """, elem_id="lut_markdown")
-                    with gr.Column():
-                        gr.Examples(elem_id="lut_examples",
-                            examples=[[f] for f in constants.lut_files],
-                            inputs=[lut_filename],
-                            outputs=[lut_filename],
-                            label="Select a Filter (LUT) file. Populate the LUT File Name field"
-                        )
-                    
-                with gr.Row():
-                    apply_lut_button = gr.Button("Apply Filter (LUT)", elem_classes="solid", elem_id="apply_lut_button")
-                    
-                lut_file.change(get_filename, inputs=[lut_file], outputs=[lut_filename])
-                lut_filename.change(show_lut, inputs=[lut_filename, lut_example_image], outputs=[lut_example_image])
-                apply_lut_button.click(
-                    lambda lut_filename, input_image: gr.Warning("Please upload an Input Image to get started.") if input_image is None else apply_lut_to_image_path(lut_filename, input_image)[0], 
-                    inputs=[lut_filename, input_image], 
-                    outputs=[input_image], 
-                    scroll_to_output=True
+                ### Why You'll Love It
+                - **Fun and Easy to Use:** With an intuitive interface and real-time previews, creating hexagon grids has never been this fun!
+                - **Endless Creativity:** Unleash your creativity with endless customization options and see your images transform in unique ways.
+                - **Hexagon-Inspired Theme:** Enjoy a delightful yellow and purple theme inspired by hexagons! ⬢
+                - **Advanced AI Models:** Leverage advanced AI models and LoRA weights for high-quality image generation and customization.
+
+                ### Get Started
+                1. **Upload or Generate an Image:** Start by uploading your own image or generate one using our AI-powered tool.
+                2. **Customize Your Grid:** Play around with the settings to create the perfect hexagon grid overlay.
+                3. **Download and Share:** Once you're happy with your creation, download it and share it with the world!
+
+                ### Advanced Features
+                - **Generative AI Integration:** Utilize models like `black-forest-labs/FLUX.1-dev` and various LoRA weights for generating unique images.
+                - **Pre-rendered Maps:** Access a library of pre-rendered hexagon maps for quick and easy customization.
+                - **Image Filter [Look-Up Table (LUT)] Application:** Apply filters (LUTs) to your images for color grading and enhancement.
+                - **Depth and 3D Model Generation:** Create depth maps and 3D models from your images for enhanced visualization.
+                - **Add Margins:** Customize margins around your images for a polished finish.
+
+                Join the hive and start creating with HexaGrid Creator today!
+        
+                """, elem_classes="intro")
+        with gr.Row():
+            from utils.image_utils import convert_to_rgba_png
+
+            # Existing code
+            with gr.Column(scale=2):
+                input_image = gr.Image(
+                    label="Input Image",
+                    type="filepath",
+                    interactive=True,
+                    elem_classes="centered solid imgcontainer",
+                    key="imgInput",
+                    image_mode=None,
+                    format="PNG", 
+                    show_download_button=True,
                 )
 
-    with gr.Row():
-        with gr.Accordion("Generative AI", open = False):
-            with gr.Row():
-                with gr.Column():
-                    model_options = gr.Dropdown(
-                        label="Model Options",
-                        choices=constants.MODELS + constants.LORA_WEIGHTS + ["Manual Entry"],
-                        value="Cossale/Frames2-Flex.1",
-                        elem_classes="solid"
-                    )
-                    model_textbox = gr.Textbox(
-                        label="LORA/Model",
-                        value="Cossale/Frames2-Flex.1",
-                        elem_classes="solid",
-                        elem_id="inference_model",
-                        visible=False
-                    )
-                    # Update map_options to a Dropdown with choices from constants.PROMPTS keys
+            # New code to convert input image to RGBA PNG
+            def on_input_image_change(image_path):
+                if image_path is None:
+                    gr.Warning("Please upload an Input Image to get started.")
+                    return None
+                img, img_path = convert_to_rgba_png(image_path)
+                return img_path
+
+            input_image.change(
+                fn=on_input_image_change,
+                inputs=[input_image],
+                outputs=[input_image], scroll_to_output=True,
+            )
+            with gr.Column():
+                with gr.Accordion("Hex Coloring and Exclusion", open = False):
                     with gr.Row():
                         with gr.Column():
-                                map_options = gr.Dropdown(
-                                label="Map Options",
-                                choices=list(constants.PROMPTS.keys()),
-                                value="Alien Landscape",
-                                elem_classes="solid",
-                                scale=0
-                            )
+                            color_picker = gr.ColorPicker(label="Pick a color to exclude",value="#505050")
+                        with gr.Column():
+                            filter_color = gr.Checkbox(label="Filter Excluded Colors from Sampling", value=False,)
+                    exclude_color_button = gr.Button("Exclude Color", elem_id="exlude_color_button", elem_classes="solid")
+                    color_display = gr.DataFrame(label="List of Excluded RGBA Colors", headers=["R", "G", "B", "A"], elem_id="excluded_colors", type="array", value=build_dataframe(excluded_color_list), interactive=True, elem_classes="solid centered")
+                    selected_row = gr.Number(0, label="Selected Row", visible=False)
+                    delete_button = gr.Button("Delete Row", elem_id="delete_exclusion_button", elem_classes="solid")
+                    fill_hex = gr.Checkbox(label="Fill Hex with color from Image", value=True)
+                with gr.Accordion("Image Filters", open = False):
+                    with gr.Row():
+                        with gr.Column():
+                            composite_color = gr.ColorPicker(label="Color", value="#ede9ac44")
+                        with gr.Column():
+                            composite_opacity = gr.Slider(label="Opacity %", minimum=0, maximum=100, value=50, interactive=True)
+                    with gr.Row():
+                        composite_button = gr.Button("Composite", elem_classes="solid")
+                    with gr.Row():
+                        with gr.Column():
+                            lut_filename = gr.Textbox(
+                                value="", 
+                                label="Look Up Table (LUT) File Name",
+                                elem_id="lutFileName")
+                        with gr.Column():
+                            lut_file = gr.File(
+                                value=None,
+                                file_count="single",
+                                file_types=[".cube"],
+                                type="filepath",
+                                label="LUT cube File")
+                    with gr.Row():
+                        lut_example_image = gr.Image(type="pil", label="Filter (LUT) Example Image", value=constants.default_lut_example_img)
+                    with gr.Row():
                         with gr.Column():
-                            # Add Dropdown for sizing of Images, height and width based on selection. Options are 16x9, 16x10, 4x5, 1x1
-                            # The values of height and width are based on common resolutions for each aspect ratio
-                            # Default to 16x9, 912x512
-                            image_size_ratio = gr.Dropdown(label="Image Size", choices=["16:9", "16:10", "4:5", "4:3", "2:1","3:2","1:1", "9:16", "10:16", "5:4", "3:4","1:2", "2:3"], value="16:9", elem_classes="solid", type="value", scale=0, interactive=True)
+                            gr.Markdown("""
+                            ### Included Filters (LUTs)
+                            There are several included Filters:
+
+                            Try them on the example image before applying to your Input Image.
+                            """, elem_id="lut_markdown")
                         with gr.Column():
-                            seed = gr.Slider(
-                                label="Seed",
-                                minimum=0,
-                                maximum=constants.MAX_SEED,
-                                step=1,
-                                value=0,
-                                scale=0
+                            gr.Examples(elem_id="lut_examples",
+                                examples=[[f] for f in constants.lut_files],
+                                inputs=[lut_filename],
+                                outputs=[lut_filename],
+                                label="Select a Filter (LUT) file. Populate the LUT File Name field"
                             )
-                            randomize_seed = gr.Checkbox(label="Randomize seed", value=True, scale=0, interactive=True)
-                    prompt_textbox = gr.Textbox(
-                        label="Prompt",
-                        visible=False,
-                        elem_classes="solid",
-                        value="top-down, (tabletop_map built from small hexagon pieces) hexagon map of a Battletech_boardgame forest with lakes, forest, magic fauna, and snow at the top and bottom, (middle is dark, no_reflections, no_shadows) , tall and short hexagon tiles. Viewed from above.",
-                        lines=4
-                    )
-                    negative_prompt_textbox = gr.Textbox(
-                        label="Negative Prompt",
-                        visible=False,
-                        elem_classes="solid",
-                        value="low quality, bad anatomy, blurry, cropped, worst quality, shadows, people, humans, reflections, shadows, realistic map of the Earth, isometric, text"
-                    )
-                    prompt_notes_label = gr.Label(
-                        "You should use FRM$ as trigger words. @1.5 minutes",
-                        elem_classes="solid centered small",
-                        show_label=False,
-                        visible=False
-                    )
-                    # Keep the change event to maintain functionality
-                    map_options.change(
-                        fn=update_prompt_visibility,
-                        inputs=[map_options],
-                        outputs=[prompt_textbox, negative_prompt_textbox, prompt_notes_label]
-                    )
+                    
                     with gr.Row():
-                        generate_input_image = gr.Button(
-                            "Generate AI Image",
-                            elem_id="generate_input_image",
-                            elem_classes="solid"
+                        apply_lut_button = gr.Button("Apply Filter (LUT)", elem_classes="solid", elem_id="apply_lut_button")
+                    
+                    lut_file.change(get_filename, inputs=[lut_file], outputs=[lut_filename])
+                    lut_filename.change(show_lut, inputs=[lut_filename, lut_example_image], outputs=[lut_example_image])
+                    apply_lut_button.click(
+                        lambda lut_filename, input_image: gr.Warning("Please upload an Input Image to get started.") if input_image is None else apply_lut_to_image_path(lut_filename, input_image)[0], 
+                        inputs=[lut_filename, input_image], 
+                        outputs=[input_image], 
+                        scroll_to_output=True
                     )
-                with gr.Column(scale=2):
-                     with gr.Accordion("Template Image Styles", open = False):
-                         with gr.Row():
-                            # 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",file_types=["image"], format="png",allow_preview=False)
-                         with gr.Row():
-                             image_guidance_stength = gr.Slider(label="Image Guidance Strength", minimum=0, maximum=1.0, value=0.5, step=0.05, interactive=True)
-                             with gr.Column():
-                                 replace_input_image_button = gr.Button(
-                                     "Replace Input Image",
-                                     elem_id="prerendered_replace_input_image_button",
-                                     elem_classes="solid"
-                                 )
-                             with gr.Column():
-                                 generate_input_image_from_gallery = gr.Button(
-                                     "Generate AI Image from Gallery",
-                                     elem_id="generate_input_image_from_gallery",
-                                     elem_classes="solid"
-                                 )
-
-    with gr.Accordion("Advanced Hexagon Settings", open = False):
-        with gr.Row():
-            start_x = gr.Number(label="Start X", value=0, minimum=-512, maximum= 512, precision=0)
-            start_y = gr.Number(label="Start Y", value=0, minimum=-512, maximum= 512, precision=0)
-            end_x = gr.Number(label="End X", value=0, minimum=-512, maximum= 512, precision=0)
-            end_y = gr.Number(label="End Y", value=0, minimum=-512, maximum= 512, precision=0)
-        with gr.Row():
-            x_spacing = gr.Number(label="Adjust Horizontal spacing", value=-1, minimum=-200, maximum=200, precision=1)
-            y_spacing = gr.Number(label="Adjust Vertical spacing", value=1, minimum=-200, maximum=200, precision=1)
-        with gr.Row():
-            rotation = gr.Slider(-90, 180, 0.0, 0.1, label="Hexagon Rotation (degree)")
-            add_hex_text = gr.Dropdown(label="Add Text to Hexagons", choices=[None, "Row-Column Coordinates", "Sequential Numbers", "Playing Cards Sequential", "Playing Cards Alternate Red and Black", "Custom List"], value=None)
-        with gr.Row():
-            custom_text_list = gr.TextArea(label="Custom Text List", value=constants.cards_alternating, visible=False,)
-            custom_text_color_list = gr.TextArea(label="Custom Text Color List", value=constants.card_colors_alternating, visible=False)
-        with gr.Row():
-            hex_text_info = gr.Markdown("""
-            ### Text Color uses the Border Color and Border Opacity, unless you use a custom list.
-            ### The Custom Text List and Custom Text Color List are comma separated lists.
-            ### The custom color list is a comma separated list of hex colors.
-            #### Example: "A,2,3,4,5,6,7,8,9,10,J,Q,K", "red,#0000FF,#00FF00,red,#FFFF00,#00FFFF,#FF8000,#FF00FF,#FF0080,#FF8000,#FF0080,lightblue"
-            """, elem_id="hex_text_info", visible=False)
-        add_hex_text.change(
-        fn=lambda x: (
-            gr.update(visible=(x == "Custom List")), 
-            gr.update(visible=(x == "Custom List")),
-            gr.update(visible=(x != None))
-            ),
-            inputs=add_hex_text,
-            outputs=[custom_text_list, custom_text_color_list, hex_text_info]
-        )
-    with gr.Row():
-        hex_size = gr.Number(label="Hexagon Size", value=32, minimum=1, maximum=768)
-        border_size = gr.Slider(-5,25,value=0,step=1,label="Border Size")
-    with gr.Row():        
-        background_color = gr.ColorPicker(label="Background Color", value="#000000", interactive=True)
-        background_opacity = gr.Slider(0,100,0,1,label="Background Opacity %")
-        border_color = gr.ColorPicker(label="Border Color", value="#7b7b7b", interactive=True)
-        border_opacity = gr.Slider(0,100,0,1,label="Border Opacity %")
-    with gr.Row():
-        hex_button = gr.Button("Generate Hex Grid!", elem_classes="solid", elem_id="btn-generate")
-    with gr.Row():
-        output_image = gr.Image(label="Hexagon Grid Image", image_mode = "RGBA", show_download_button=True, show_share_button=True,elem_classes="centered solid imgcontainer", format="PNG", type="filepath", key="ImgOutput")
-        overlay_image = gr.Image(label="Hexagon Overlay Image", image_mode = "RGBA", show_share_button=True, elem_classes="centered solid imgcontainer", format="PNG", type="filepath", key="ImgOverlay")
-    with gr.Row():
-        output_overlay_composite = gr.Slider(0,100,50,0.5, label="Interpolate Intensity")
-        output_blend_multiply_composite = gr.Slider(0,100,50,0.5, label="Overlay Intensity")
-        output_alpha_composite = gr.Slider(0,100,50,0.5, label="Alpha Composite Intensity")
-    with gr.Accordion("Add Margins (bleed)", open=False):
+
         with gr.Row():
-            border_image_source = gr.Radio(label="Add Margins around which Image", choices=["Input Image", "Overlay Image"], value="Overlay Image")
+            with gr.Accordion("Generative AI", open = False):
+                with gr.Row():
+                    with gr.Column():
+                        model_options = gr.Dropdown(
+                            label="Model Options",
+                            choices=constants.MODELS + constants.LORA_WEIGHTS + ["Manual Entry"],
+                            value="Cossale/Frames2-Flex.1",
+                            elem_classes="solid"
+                        )
+                        model_textbox = gr.Textbox(
+                            label="LORA/Model",
+                            value="Cossale/Frames2-Flex.1",
+                            elem_classes="solid",
+                            elem_id="inference_model",
+                            visible=False
+                        )
+                        # Update map_options to a Dropdown with choices from constants.PROMPTS keys
+                        with gr.Row():
+                            with gr.Column():
+                                    map_options = gr.Dropdown(
+                                    label="Map Options",
+                                    choices=list(constants.PROMPTS.keys()),
+                                    value="Alien Landscape",
+                                    elem_classes="solid",
+                                    scale=0
+                                )
+                            with gr.Column():
+                                # Add Dropdown for sizing of Images, height and width based on selection. Options are 16x9, 16x10, 4x5, 1x1
+                                # The values of height and width are based on common resolutions for each aspect ratio
+                                # Default to 16x9, 912x512
+                                image_size_ratio = gr.Dropdown(label="Image Size", choices=["16:9", "16:10", "4:5", "4:3", "2:1","3:2","1:1", "9:16", "10:16", "5:4", "3:4","1:2", "2:3"], value="16:9", elem_classes="solid", type="value", scale=0, interactive=True)
+                            with gr.Column():
+                                seed_slider = gr.Slider(
+                                    label="Seed",
+                                    minimum=0,
+                                    maximum=constants.MAX_SEED,
+                                    step=1,
+                                    value=0,
+                                    scale=0
+                                )
+                                randomize_seed = gr.Checkbox(label="Randomize seed", value=True, scale=0, interactive=True)
+                        prompt_textbox = gr.Textbox(
+                            label="Prompt",
+                            visible=False,
+                            elem_classes="solid",
+                            value="top-down, (rectangular tabletop_map) alien planet map, Battletech_boardgame scifi world with forests, lakes, oceans, continents and snow at the top and bottom, (middle is dark, no_reflections, no_shadows), from directly above. From 100,000 feet looking straight down",
+                            lines=4
+                        )
+                        negative_prompt_textbox = gr.Textbox(
+                            label="Negative Prompt",
+                            visible=False,
+                            elem_classes="solid",
+                            value="Earth, low quality, bad anatomy, blurry, cropped, worst quality, shadows, people, humans, reflections, shadows, realistic map of the Earth, isometric, text"
+                        )
+                        prompt_notes_label = gr.Label(
+                            "You should use FRM$ as trigger words. @1.5 minutes",
+                            elem_classes="solid centered small",
+                            show_label=False,
+                            visible=False
+                        )
+                        # Keep the change event to maintain functionality
+                        map_options.change(
+                            fn=update_prompt_visibility,
+                            inputs=[map_options],
+                            outputs=[prompt_textbox, negative_prompt_textbox, prompt_notes_label]
+                        )
+                        with gr.Row():
+                            generate_input_image = gr.Button(
+                                "Generate AI Image",
+                                elem_id="generate_input_image",
+                                elem_classes="solid"
+                        )
+                    with gr.Column(scale=2):
+                         with gr.Accordion("Template Image Styles", open = False):
+                             with gr.Row():
+                                # 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)
+                                 with gr.Column():
+                                     replace_input_image_button = gr.Button(
+                                         "Replace Input Image",
+                                         elem_id="prerendered_replace_input_image_button",
+                                         elem_classes="solid"
+                                     )
+                                 with gr.Column():
+                                     generate_input_image_from_gallery = gr.Button(
+                                         "Generate AI Image from Gallery",
+                                         elem_id="generate_input_image_from_gallery",
+                                         elem_classes="solid"
+                                     )
+
+        with gr.Accordion("Advanced Hexagon Settings", open = False):
+            with gr.Row():
+                start_x = gr.Number(label="Start X", value=0, minimum=-512, maximum= 512, precision=0)
+                start_y = gr.Number(label="Start Y", value=0, minimum=-512, maximum= 512, precision=0)
+                end_x = gr.Number(label="End X", value=0, minimum=-512, maximum= 512, precision=0)
+                end_y = gr.Number(label="End Y", value=0, minimum=-512, maximum= 512, precision=0)
+            with gr.Row():
+                x_spacing = gr.Number(label="Adjust Horizontal spacing", value=-1, minimum=-200, maximum=200, precision=1)
+                y_spacing = gr.Number(label="Adjust Vertical spacing", value=1, minimum=-200, maximum=200, precision=1)
+            with gr.Row():
+                rotation = gr.Slider(-90, 180, 0.0, 0.1, label="Hexagon Rotation (degree)")
+                add_hex_text = gr.Dropdown(label="Add Text to Hexagons", choices=[None, "Row-Column Coordinates", "Sequential Numbers", "Playing Cards Sequential", "Playing Cards Alternate Red and Black", "Custom List"], value=None)
+            with gr.Row():
+                custom_text_list = gr.TextArea(label="Custom Text List", value=constants.cards_alternating, visible=False,)
+                custom_text_color_list = gr.TextArea(label="Custom Text Color List", value=constants.card_colors_alternating, visible=False)
+            with gr.Row():
+                hex_text_info = gr.Markdown("""
+                ### Text Color uses the Border Color and Border Opacity, unless you use a custom list.
+                ### The Custom Text List and Custom Text Color List are comma separated lists.
+                ### The custom color list is a comma separated list of hex colors.
+                #### Example: "A,2,3,4,5,6,7,8,9,10,J,Q,K", "red,#0000FF,#00FF00,red,#FFFF00,#00FFFF,#FF8000,#FF00FF,#FF0080,#FF8000,#FF0080,lightblue"
+                """, elem_id="hex_text_info", visible=False)
+            add_hex_text.change(
+            fn=lambda x: (
+                gr.update(visible=(x == "Custom List")), 
+                gr.update(visible=(x == "Custom List")),
+                gr.update(visible=(x != None))
+                ),
+                inputs=add_hex_text,
+                outputs=[custom_text_list, custom_text_color_list, hex_text_info]
+            )
         with gr.Row():
-            mask_width = gr.Number(label="Margins Width", value=10, minimum=0, maximum=100, precision=0)
-            mask_height = gr.Number(label="Margins Height", value=10, minimum=0, maximum=100, precision=0)
+            hex_size = gr.Number(label="Hexagon Size", value=32, minimum=1, maximum=768)
+            border_size = gr.Slider(-5,25,value=0,step=1,label="Border Size")
+        with gr.Row():        
+            background_color = gr.ColorPicker(label="Background Color", value="#000000", interactive=True)
+            background_opacity = gr.Slider(0,100,0,1,label="Background Opacity %")
+            border_color = gr.ColorPicker(label="Border Color", value="#7b7b7b", interactive=True)
+            border_opacity = gr.Slider(0,100,0,1,label="Border Opacity %")
         with gr.Row():
-                margin_color = gr.ColorPicker(label="Margin Color", value="#333333FF", interactive=True)
-                margin_opacity = gr.Slider(0,100,95,0.5,label="Margin Opacity %")
+            hex_button = gr.Button("Generate Hex Grid!", elem_classes="solid", elem_id="btn-generate")
         with gr.Row():
-            add_border_button = gr.Button("Add Margins", elem_classes="solid", variant="secondary")
+            output_image = gr.Image(label="Hexagon Grid Image", image_mode = "RGBA", show_download_button=True, show_share_button=True,elem_classes="centered solid imgcontainer", format="PNG", type="filepath", key="ImgOutput")
+            overlay_image = gr.Image(label="Hexagon Overlay Image", image_mode = "RGBA", show_share_button=True, elem_classes="centered solid imgcontainer", format="PNG", type="filepath", key="ImgOverlay")
         with gr.Row():
-            bordered_image_output = gr.Image(label="Image with Margins", image_mode="RGBA", show_download_button=True, show_share_button=True, elem_classes="centered solid imgcontainer", format="PNG", type="filepath", key="ImgBordered")
+            output_overlay_composite = gr.Slider(0,100,50,0.5, label="Interpolate Intensity")
+            output_blend_multiply_composite = gr.Slider(0,100,50,0.5, label="Overlay Intensity")
+            output_alpha_composite = gr.Slider(0,100,50,0.5, label="Alpha Composite Intensity")
+        with gr.Accordion("Add Margins (bleed)", open=False):
+            with gr.Row():
+                border_image_source = gr.Radio(label="Add Margins around which Image", choices=["Input Image", "Overlay Image"], value="Overlay Image")
+            with gr.Row():
+                mask_width = gr.Number(label="Margins Width", value=10, minimum=0, maximum=100, precision=0)
+                mask_height = gr.Number(label="Margins Height", value=10, minimum=0, maximum=100, precision=0)
+            with gr.Row():
+                    margin_color = gr.ColorPicker(label="Margin Color", value="#333333FF", interactive=True)
+                    margin_opacity = gr.Slider(0,100,95,0.5,label="Margin Opacity %")
+            with gr.Row():
+                add_border_button = gr.Button("Add Margins", elem_classes="solid", variant="secondary")
+            with gr.Row():
+                bordered_image_output = gr.Image(label="Image with Margins", image_mode="RGBA", show_download_button=True, show_share_button=True, elem_classes="centered solid imgcontainer", format="PNG", type="filepath", key="ImgBordered")
 
-    with gr.Accordion("Height Maps and 3D", open = False):
-        with gr.Row():
-            with gr.Column():
-                voxel_size_factor = gr.Slider(label="Voxel Size Factor", value=1.00, minimum=0.01, maximum=40.00, step=0.01)
-            with gr.Column():
-                depth_image_source = gr.Radio(label="Depth Image Source", choices=["Input Image", "Output Image", "Overlay Image","Image with Margins"], value="Input Image")
+        with gr.Accordion("Height Maps and 3D", open = False):
+            with gr.Row():
+                with gr.Column():
+                    voxel_size_factor = gr.Slider(label="Voxel Size Factor", value=1.00, minimum=0.01, maximum=40.00, step=0.01)
+                with gr.Column():
+                    depth_image_source = gr.Radio(label="Depth Image Source", choices=["Input Image", "Output Image", "Overlay Image","Image with Margins"], value="Input Image")
+            with gr.Row():
+                generate_depth_button = gr.Button("Generate Depth Map and 3D Model From Selected Image", elem_classes="solid", variant="secondary")
+            with gr.Row():
+                depth_map_output = gr.Image(label="Depth Map", image_mode="L", elem_classes="centered solid imgcontainer", format="PNG", type="filepath", key="ImgDepth")
+                model_output = gr.Model3D(label="3D Model", clear_color=[1.0, 1.0, 1.0, 0.25], key="Img3D", elem_classes="centered solid imgcontainer")
         with gr.Row():
-            generate_depth_button = gr.Button("Generate Depth Map and 3D Model From Selected Image", elem_classes="solid", variant="secondary")
+            gr.Examples(examples=[
+                ["assets//examples//hex_map_p1.png", False, True, -32,-31,80,80,-1.8,0,35,0,1,"#FFD0D0", 15],
+                ["assets//examples//hex_map_p1_overlayed.png", False, False, -32,-31,80,80,-1.8,0,35,0,1,"#FFD0D0", 75],
+                ["assets//examples//hex_flower_logo.png", False, True, -95,-95,100,100,-24,-2,190,30,2,"#FF8951", 50],
+                ["assets//examples//hexed_fract_1.png", False, True, 0,0,0,0,0,0,10,0,0,"#000000", 5],
+                ["assets//examples//tmpzt3mblvk.png", False, True, -20,10,0,0,-6,-2,35,30,1,"#ffffff", 0],
+                ],
+                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():
-            depth_map_output = gr.Image(label="Depth Map", image_mode="L", elem_classes="centered solid imgcontainer", format="PNG", type="filepath", key="ImgDepth")
-            model_output = gr.Model3D(label="3D Model", clear_color=[1.0, 1.0, 1.0, 0.25], key="Img3D", elem_classes="centered solid imgcontainer")
-    with gr.Row():
-        gr.Examples(examples=[
-            ["assets//examples//hex_map_p1.png", False, True, -32,-31,80,80,-1.8,0,35,0,1,"#FFD0D0", 15],
-            ["assets//examples//hex_map_p1_overlayed.png", False, False, -32,-31,80,80,-1.8,0,35,0,1,"#FFD0D0", 75],
-            ["assets//examples//hex_flower_logo.png", False, True, -95,-95,100,100,-24,-2,190,30,2,"#FF8951", 50],
-            ["assets//examples//hexed_fract_1.png", False, True, 0,0,0,0,0,0,10,0,0,"#000000", 5],
-            ["assets//examples//tmpzt3mblvk.png", False, True, -20,10,0,0,-6,-2,35,30,1,"#ffffff", 0],
-            ],
-            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")
-
-    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)])
-
-    delete_button.click(fn=delete_color, inputs=[selected_row, color_display], outputs=[color_display])
-    exclude_color_button.click(fn=add_color, inputs=[color_picker, gr.State(excluded_color_list)], outputs=[color_display, gr.State(excluded_color_list)])
-    hex_button.click(
-        fn=lambda hex_size, border_size, input_image, start_x, start_y, end_x, end_y, rotation, background_color, background_opacity, border_color, border_opacity, fill_hex, color_display, filter_color, x_spacing, y_spacing, add_hex_text, custom_text_list, custom_text_color_list: 
-            gr.Warning("Please upload an Input Image to get started.") if input_image is None else hex_create(hex_size, border_size, input_image, start_x, start_y, end_x, end_y, rotation, background_color, background_opacity, border_color, border_opacity, fill_hex, color_display, filter_color, x_spacing, y_spacing, add_hex_text, custom_text_list, custom_text_color_list),
-        inputs=[hex_size, border_size, input_image, start_x, start_y, end_x, end_y, rotation, background_color, background_opacity, border_color, border_opacity, fill_hex, color_display, filter_color, x_spacing, y_spacing, add_hex_text, custom_text_list, custom_text_color_list],
-        outputs=[output_image, overlay_image],
-        scroll_to_output=True
-    )
-    generate_input_image.click(
-        fn=generate_input_image_click,
-        inputs=[map_options, prompt_textbox, negative_prompt_textbox, model_textbox,gr.State( seed if randomize_seed==False else random.randint(0, constants.MAX_SEED)), gr.State(False), gr.State(0.5), image_size_ratio],
-        outputs=[input_image], scroll_to_output=True
-    )
-    generate_depth_button.click(
-        fn=generate_depth_button_click,
-        inputs=[depth_image_source, voxel_size_factor, input_image, output_image, overlay_image, bordered_image_output],
-        outputs=[depth_map_output, model_output], scroll_to_output=True
-    )
-    model_textbox.change(
-        fn=update_prompt_notes,
-        inputs=model_textbox,
-        outputs=prompt_notes_label,preprocess=False
-    )
-    model_options.change(
-        fn=lambda x: (gr.update(visible=(x == "Manual Entry")), gr.update(value=x) if x != "Manual Entry" else gr.update()),
-        inputs=model_options,
-        outputs=[model_textbox, model_textbox]
-    )
-    model_options.change(
-        fn=update_prompt_notes,
-        inputs=model_options,
-        outputs=prompt_notes_label
-    )
-    composite_button.click(
-        fn=lambda input_image, composite_color, composite_opacity: gr.Warning("Please upload an Input Image to get started.") if input_image is None else change_color(input_image, composite_color, composite_opacity),
-        inputs=[input_image, composite_color, composite_opacity],
-        outputs=[input_image]
-    )
+            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")
 
-    #use conditioned_image as the input_image for generate_input_image_click
-    generate_input_image_from_gallery.click(
-        fn=generate_input_image_click,
-        inputs=[map_options, prompt_textbox, negative_prompt_textbox, model_textbox, gr.State(True), image_guidance_stength, image_size_ratio],
-        outputs=[input_image], scroll_to_output=True
-    )
+        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=[]
+        )
 
-    # Update the state variable with the prerendered image filepath when an image is selected
-    prerendered_image_gallery.select(
-        fn=on_prerendered_gallery_selection, 
-        inputs=None, 
-        outputs=[gr.State(current_prerendered_image)],  # Update the state with the selected image
-        show_api=False
-    )
-    # replace input image with selected gallery image
-    replace_input_image_button.click(
-        lambda: current_prerendered_image.value,
-        inputs=None,
-        outputs=[input_image], scroll_to_output=True
-    )
-    output_overlay_composite.change(
-        fn=combine_images_with_lerp,
-        inputs=[input_image, output_image, output_overlay_composite],
-        outputs=[overlay_image], scroll_to_output=True
-    )
-    output_blend_multiply_composite.change(
-        fn=multiply_and_blend_images,
-        inputs=[input_image, output_image, output_blend_multiply_composite],
-        outputs=[overlay_image],
-        scroll_to_output=True
-    )
-    output_alpha_composite.change(
-        fn=alpha_composite_with_control,
-        inputs=[input_image, output_image, output_alpha_composite],
-        outputs=[overlay_image],
-        scroll_to_output=True
-    )
-    add_border_button.click(
-        fn=lambda image_source, mask_w, mask_h, color, opacity, input_img, overlay_img: add_border(input_img if image_source == "Input Image" else overlay_img, mask_w, mask_h, update_color_opacity(detect_color_format(color), opacity * 2.55)),
-        inputs=[border_image_source, mask_width, mask_height, margin_color, margin_opacity, input_image, overlay_image],
-        outputs=[bordered_image_output],
-        scroll_to_output=True
-    )
-(())
-if __name__ == "__main__":
+        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)])
+
+        delete_button.click(fn=delete_color, inputs=[selected_row, color_display], outputs=[color_display])
+        exclude_color_button.click(fn=add_color, inputs=[color_picker, gr.State(excluded_color_list)], outputs=[color_display, gr.State(excluded_color_list)])
+        hex_button.click(
+            fn=lambda hex_size, border_size, input_image, start_x, start_y, end_x, end_y, rotation, background_color, background_opacity, border_color, border_opacity, fill_hex, color_display, filter_color, x_spacing, y_spacing, add_hex_text, custom_text_list, custom_text_color_list: 
+                gr.Warning("Please upload an Input Image to get started.") if input_image is None else hex_create(hex_size, border_size, input_image, start_x, start_y, end_x, end_y, rotation, background_color, background_opacity, border_color, border_opacity, fill_hex, color_display, filter_color, x_spacing, y_spacing, add_hex_text, custom_text_list, custom_text_color_list),
+            inputs=[hex_size, border_size, input_image, start_x, start_y, end_x, end_y, rotation, background_color, background_opacity, border_color, border_opacity, fill_hex, color_display, filter_color, x_spacing, y_spacing, add_hex_text, custom_text_list, custom_text_color_list],
+            outputs=[output_image, overlay_image],
+            scroll_to_output=True
+        )
+        generate_input_image.click(
+            fn=generate_input_image_click,
+            inputs=[map_options, prompt_textbox, negative_prompt_textbox, model_textbox, randomize_seed, seed_slider, gr.State(False), gr.State(0.5), image_size_ratio],
+            outputs=[input_image], scroll_to_output=True
+        )
+        generate_depth_button.click(
+            fn=generate_depth_button_click,
+            inputs=[depth_image_source, voxel_size_factor, input_image, output_image, overlay_image, bordered_image_output],
+            outputs=[depth_map_output, model_output], scroll_to_output=True
+        )
+        model_textbox.change(
+            fn=update_prompt_notes,
+            inputs=model_textbox,
+            outputs=prompt_notes_label,preprocess=False
+        )
+        model_options.change(
+            fn=lambda x: (gr.update(visible=(x == "Manual Entry")), gr.update(value=x) if x != "Manual Entry" else gr.update()),
+            inputs=model_options,
+            outputs=[model_textbox, model_textbox]
+        )
+        model_options.change(
+            fn=update_prompt_notes,
+            inputs=model_options,
+            outputs=prompt_notes_label
+        )
+        composite_button.click(
+            fn=lambda input_image, composite_color, composite_opacity: gr.Warning("Please upload an Input Image to get started.") if input_image is None else change_color(input_image, composite_color, composite_opacity),
+            inputs=[input_image, composite_color, composite_opacity],
+            outputs=[input_image]
+        )
+
+        #use conditioned_image as the input_image for generate_input_image_click
+        generate_input_image_from_gallery.click(
+            fn=generate_input_image_click,
+            inputs=[map_options, prompt_textbox, negative_prompt_textbox, model_textbox,randomize_seed, seed_slider, gr.State(True), image_guidance_stength, image_size_ratio],
+            outputs=[input_image], scroll_to_output=True
+        )
+
+        # Update the state variable with the prerendered image filepath when an image is selected
+        prerendered_image_gallery.select(
+            fn=on_prerendered_gallery_selection, 
+            inputs=None, 
+            outputs=[gr.State(current_prerendered_image)],  # Update the state with the selected image
+            show_api=False
+        )
+        # replace input image with selected gallery image
+        replace_input_image_button.click(
+            lambda: current_prerendered_image.value,
+            inputs=None,
+            outputs=[input_image], scroll_to_output=True
+        )
+        output_overlay_composite.change(
+            fn=combine_images_with_lerp,
+            inputs=[input_image, output_image, output_overlay_composite],
+            outputs=[overlay_image], scroll_to_output=True
+        )
+        output_blend_multiply_composite.change(
+            fn=multiply_and_blend_images,
+            inputs=[input_image, output_image, output_blend_multiply_composite],
+            outputs=[overlay_image],
+            scroll_to_output=True
+        )
+        output_alpha_composite.change(
+            fn=alpha_composite_with_control,
+            inputs=[input_image, output_image, output_alpha_composite],
+            outputs=[overlay_image],
+            scroll_to_output=True
+        )
+        add_border_button.click(
+            fn=lambda image_source, mask_w, mask_h, color, opacity, input_img, overlay_img: add_border(input_img if image_source == "Input Image" else overlay_img, mask_w, mask_h, update_color_opacity(detect_color_format(color), opacity * 2.55)),
+            inputs=[border_image_source, mask_width, mask_height, margin_color, margin_opacity, input_image, overlay_image],
+            outputs=[bordered_image_output],
+            scroll_to_output=True
+        )
+    (())
     beeuty.queue(default_concurrency_limit=1,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")
+    beeuty.launch(allowed_paths=["assets","/","./assets","images","./images", "./images/prerendered"], favicon_path="./assets/favicon.ico", max_file_size="10mb")
+
+
+if __name__ == "__main__":
+    logging.basicConfig(
+        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")

pre-requirements.txt

@@ -1 +1,5 @@
-pip>=25.0
+pip>=25.0.1
+#torch==2.4.0 --extra-index-url https://download.pytorch.org/whl/cu124
+#torch-2.4.0%2Bcu124-cp310-cp310-linux_x86_64.whl#sha256=2cb28155635e3d3d0be198e3f3e7457a1d7b99e8c2eedc73fe22fab574d11a4c
+#torchvision==0.19.0 --extra-index-url https://download.pytorch.org/whl/cu124
+#/torchvision-0.19.0%2Bcu124-cp310-cp310-linux_x86_64.whl#sha256=82cf10450537aeb9584ceaf53633f177bb809d563c5d64526f4b9be7668b2769

requirements.txt

@@ -1,37 +1,44 @@
+git+https://github.com/huggingface/diffusers.git
+#diffusers[torch]==0.32.2
+transformers
 accelerate
+safetensors
+sentencepiece
 invisible_watermark
+
 # Updated versions 2.4.0+cu124
-#torch==2.4.0 --index-url https://download.pytorch.org/whl/cu124/torch-2.4.0%2Bcu124-cp310-cp310-linux_x86_64.whl#sha256=2cb28155635e3d3d0be198e3f3e7457a1d7b99e8c2eedc73fe22fab574d11a4c
-#torchvision --index-url https://download.pytorch.org/whl/cu124
-#torchaudio --index-url https://download.pytorch.org/whl/cu124
+#--extra-index-url https://download.pytorch.org/whl/cu124
+#torch==2.6.0 --index-url https://download.pytorch.org/whl/cu124/torch-2.4.0%2Bcu124-cp310-cp310-linux_x86_64.whl#sha256=2cb28155635e3d3d0be198e3f3e7457a1d7b99e8c2eedc73fe22fab574d11a4c 
+#torchvision==0.21.0 --index-url https://download.pytorch.org/whl/cu124/torchvision-0.19.0%2Bcu124-cp310-cp310-linux_x86_64.whl#sha256=82cf10450537aeb9584ceaf53633f177bb809d563c5d64526f4b9be7668b2769
+#torchaudio==2.6.0 --index-url https://download.pytorch.org/whl/cu124
 #xformers --index-url https://download.pytorch.org/whl/cu124
 # ==0.0.27.post2 --index-url https://download.pytorch.org/whl/cu118/xformers-0.0.27.post2%2Bcu118-cp310-cp310-manylinux2014_x86_64.whl#sha256=b3cdeeb9eae4547805ab8c3c645ac2fa9c6da85b46c039d9befa117e9f6f22fe
 
 #generic Torch versions
+#--extra-index-url https://download.pytorch.org/whl/cu124
 torch
-xformers
+torchvision
+#xformers #==0.0.29.post3
 
 # Other dependencies
 Haishoku
 pybind11>=2.12
 huggingface_hub
-# git+https://github.com/huggingface/[email protected].1#egg=transformers
-transformers==4.48.1
-gradio[oauth]
+# git+https://github.com/huggingface/[email protected].3#egg=transformers
+#gradio[oauth]
 Pillow>=11.0.0
 numpy
 requests
-git+https://github.com/huggingface/diffusers
-#diffusers[torch]
+
 peft
 opencv-python
 open3d
-protobuf
+protobuf #==3.20.3
 safetensors
 sentencepiece
 git+https://github.com/asomoza/image_gen_aux.git
 #git+https://github.com/huggingface/optimum.git
-#git+https://github.com/triton-lang/triton.git -not windows supported --disable in environment variable
+#git+https://github.com/triton-lang/triton.git #-not windows supported --disable in environment variable
 tiktoken
 #pilmoji[requests]==2.0.4
 #emoji==2.2.0
@@ -41,4 +48,5 @@ pangocffi
 pangocairocffi
 #tensorflow
 cairosvg
-python-dotenv
+python-dotenv
+#git+https://github.com/gradio-app/[email protected]#egg=gradio

src/condition.py

@@ -1,4 +1,6 @@
-import torch
+import spaces
+import gradio as gr
+from torch import Tensor, ones_like
 from typing import Optional, Union, List, Tuple
 from diffusers.pipelines import FluxPipeline
 from PIL import Image, ImageFilter
@@ -13,13 +15,13 @@ condition_dict = {
     "deblurring": 7,
     "fill": 9,
 }
-
+@spaces.GPU(progress=gr.Progress(track_tqdm=True))
 class Condition(object):
     def __init__(
         self,
         condition_type: str,
-        raw_img: Union[Image.Image, torch.Tensor] = None,
-        condition: Union[Image.Image, torch.Tensor] = None,
+        raw_img: Union[Image.Image, Tensor] = None,
+        condition: Union[Image.Image,Tensor] = None,
         mask=None,
     ) -> None:
         self.condition_type = condition_type
@@ -31,8 +33,8 @@ class Condition(object):
         # 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]:
+        self, condition_type: str, raw_img: Union[Image.Image, Tensor]
+    ) -> Union[Image.Image, Tensor]:
         """
         Returns the condition image.
         """
@@ -77,7 +79,7 @@ class Condition(object):
         Returns the type id of the condition.
         """
         return condition_dict[condition_type]
-    def _encode_image(self, pipe: FluxPipeline, cond_img: Image.Image) -> torch.Tensor:
+    def _encode_image(self, pipe: FluxPipeline, cond_img: Image.Image) -> Tensor:
         """
         Encodes an image condition into tokens using the pipeline.
         """
@@ -96,7 +98,7 @@ class Condition(object):
             pipe.dtype,
         )
         return cond_tokens, cond_ids
-    def encode(self, pipe: FluxPipeline) -> Tuple[torch.Tensor, torch.Tensor, int]:
+    def encode(self, pipe: FluxPipeline) -> Tuple[Tensor, Tensor, int]:
         """
         Encodes the condition into tokens, ids and type_id.
         """
@@ -113,5 +115,5 @@ class Condition(object):
             raise NotImplementedError(
                 f"Condition type {self.condition_type} not implemented"
             )
-        type_id = torch.ones_like(ids[:, :1]) * self.type_id
+        type_id = ones_like(ids[:, :1]) * self.type_id
         return tokens, ids, type_id

utils/ai_generator.py

@@ -2,11 +2,11 @@
 import gradio as gr
 import os
 import time
-from turtle import width  # Added for implementing delays
-import torch
+#from turtle import width  # Added for implementing delays
+from torch import cuda
 import random
 from utils.ai_generator_diffusers_flux import generate_ai_image_local
-from pathlib import Path
+#from pathlib import Path
 from huggingface_hub import InferenceClient
 import requests
 import io
@@ -38,12 +38,14 @@ def generate_ai_image(
     width=912,
     height=512,
     strength=0.5,
-    seed = random.randint(0, constants.MAX_SEED),
+    seed = 0,
     progress=gr.Progress(track_tqdm=True),
     *args,
     **kwargs
-):    
-    if (torch.cuda.is_available() and torch.cuda.device_count() >= 1):
+):   
+    if seed == 0:
+        seed = random.randint(0, constants.MAX_SEED)
+    if (cuda.is_available() and cuda.device_count() >= 1): # Check if a local GPU is available 
         print("Local GPU available. Generating image locally.")
         if conditioned_image is not None:
             pipeline = "FluxImg2ImgPipeline"

utils/ai_generator_diffusers_flux.py

@@ -1,19 +1,19 @@
 # utils/ai_generator_diffusers_flux.py
-import spaces
-import gradio as gr
 import os
 import utils.constants as constants
-import torch
+#import spaces
+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
-from diffusers.utils import load_image
-from huggingface_hub import hf_hub_download
+#import xformers
+#from diffusers.utils import load_image
+#from huggingface_hub import hf_hub_download
 from PIL import Image
 from tempfile import NamedTemporaryFile
-from src.condition import Condition
+
 from utils.image_utils import (
      crop_and_resize_image,
 )
@@ -21,23 +21,26 @@ from utils.version_info import (
     get_torch_info,
     get_diffusers_version,
     get_transformers_version,
-    get_xformers_version
+    get_xformers_version,
+    initialize_cuda,
+    release_torch_resources
 )
+import gc
 from utils.lora_details import get_trigger_words, approximate_token_count, split_prompt_precisely
-from utils.color_utils import detect_color_format
-import utils.misc as misc
-from pathlib import Path
+#from utils.color_utils import detect_color_format
+#import utils.misc as misc
+#from pathlib import Path
 import warnings
 warnings.filterwarnings("ignore", message=".*Torch was not compiled with flash attention.*")
-#print(torch.__version__)  # Ensure it's 2.0 or newer
-#print(torch.cuda.is_available())  # Ensure CUDA is available
+#print(torch_version_)  # Ensure it's 2.0 or newer
+#print(cuda.is_available())  # Ensure CUDA is available
 
 PIPELINE_CLASSES = {
     "FluxPipeline": FluxPipeline,
     "FluxImg2ImgPipeline": FluxImg2ImgPipeline,
     "FluxControlPipeline": FluxControlPipeline
 }
-@spaces.GPU(duration=140)
+#@spaces.GPU()
 def generate_image_from_text(
     text,
     model_name="black-forest-labs/FLUX.1-dev",
@@ -51,13 +54,14 @@ def generate_image_from_text(
     additional_parameters=None,
     progress=gr.Progress(track_tqdm=True)
 ):
-    device = "cuda" if torch.cuda.is_available() else "cpu"
+    from src.condition import Condition
+    device = "cuda" if cuda.is_available() else "cpu"
     print(f"device:{device}\nmodel_name:{model_name}\n")
 
     # Initialize the pipeline
     pipe = FluxPipeline.from_pretrained(
         model_name,
-        torch_dtype=torch.bfloat16 if device == "cuda" else torch.float32
+        torch_dtype=bfloat16 if device == "cuda" else float32
     ).to(device)
     pipe.enable_model_cpu_offload()
 
@@ -88,7 +92,7 @@ def generate_image_from_text(
                 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)
+    generator = Generator(device=device).manual_seed(seed)
     conditions = []
 
     # Handle conditioned image if provided
@@ -122,12 +126,12 @@ def generate_image_from_text(
     del conditions
     del generator
     del pipe
-    torch.cuda.empty_cache()
-    torch.cuda.ipc_collect()
+    cuda.empty_cache()
+    cuda.ipc_collect()
 
     return image
 
-@spaces.GPU(duration=140)
+#@spaces.GPU(progress=gr.Progress(track_tqdm=True))
 def generate_image_lowmem(
     text,
     neg_prompt=None,
@@ -141,50 +145,59 @@ def generate_image_lowmem(
     seed=0,
     true_cfg_scale=1.0,
     pipeline_name="FluxPipeline",
-    strength=0.75,    
+    strength=0.75,
     additional_parameters=None,
     progress=gr.Progress(track_tqdm=True)
-):    
+):  
     # 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())}")
 
-    device = "cuda" if torch.cuda.is_available() else "cpu"
+    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 torch.no_grad():
+    with 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.bfloat32
+            torch_dtype=bfloat16 if device == "cuda" else float32
         ).to(device)
         # Optionally, don't use CPU offload if not necessary
-        pipe.enable_model_cpu_offload()
+        
         # 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)
+            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()
+        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()
+            #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")
-        if pipeline_name == "FluxPipeline":
-            pipe.enable_vae_tiling()
+
         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])        
@@ -275,7 +288,7 @@ def generate_image_lowmem(
                 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)
+        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)
@@ -327,9 +340,10 @@ def generate_image_lowmem(
         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))
+    cuda.empty_cache()
+    cuda.ipc_collect()
+    print(cuda.memory_summary(device=None, abbreviated=False))
+    
     return image
 
 def generate_ai_image_local (
@@ -348,6 +362,8 @@ def generate_ai_image_local (
     strength=0.75,
     progress=gr.Progress(track_tqdm=True)
 ):
+    release_torch_resources()
+    print(f"Generating image with lowmem")
     try:
         if map_option != "Prompt":
             prompt = constants.PROMPTS[map_option]
@@ -387,6 +403,7 @@ def generate_ai_image_local (
         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,
@@ -400,15 +417,18 @@ def generate_ai_image_local (
             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}")
+            gc.collect()
             return tmp.name
     except Exception as e:
         print(f"Error generating AI image: {e}")
+        gc.collect()
         return None
 
 # does not work
@@ -419,7 +439,7 @@ def merge_LoRA_weights(model="black-forest-labs/FLUX.1-dev",
     if model_suffix not in lora_weights:
         raise ValueError(f"The model suffix '{model_suffix}' must be in the lora_weights string '{lora_weights}' to proceed.")
     
-    pipe = FluxPipeline.from_pretrained(model, torch_dtype=torch.bfloat16)
+    pipe = FluxPipeline.from_pretrained(model, torch_dtype=bfloat16)
     pipe.load_lora_weights(lora_weights)
     pipe.save_lora_weights(os.getenv("TMPDIR"))
     lora_name = lora_weights.split("/")[-1] + "-merged"

utils/constants.py

@@ -10,12 +10,18 @@ import numpy as np
 os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:256,expandable_segments:True"
 os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0'
 #os.environ["TF_CPP_MIN_LOG_LEVEL"] = '2'
+os.environ['CUDA_MODULE_LOADING']='LAZY'
 os.environ['USE_FLASH_ATTENTION'] = '1'
-#os.environ['XFORMERS_FORCE_DISABLE_TRITON']= '1'
+os.environ['XFORMERS_FORCE_DISABLE_TRITON']= '1'
 #os.environ['XFORMERS_FORCE_DISABLE_TORCHSCRIPT']= '1'
-os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
+#os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
 os.environ["PYTORCH_NVML_BASED_CUDA_CHECK"] = "1"
 os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+os.environ["NVIDIA_VISIBLE_DEVICES"] = "0"
+os.environ["ZEROGPU_VERSION"] = "2"
+os.environ["ZEROGPU_V2"] = "true"
+os.environ["ZERO_GPU_V2"] = "true"
+os.environ["ZERO_GPU_PATCH_TORCH_DEVICE"]='1'
 
 IS_SHARED_SPACE = "Surn/HexaGrid" in os.environ.get('SPACE_ID', '')
 
@@ -36,14 +42,14 @@ default_lut_example_img = "./LUT/daisy.jpg"
 MAX_SEED = np.iinfo(np.int32).max
 
 PROMPTS = {
-    "BorderBlack": "eight_color (tabletop_map built from small hexagon pieces) as ((empty black on all sides), barren alien_world_map), with light_blue_is_rivers and brown_is_mountains and red_is_volcano and [white_is_snow at the top and bottom of map] as (four_color background: light_blue, green, tan, brown), horizontal_gradient is (brown to tan to green to light_blue to blue) and vertical_gradient is (white to blue to (green, tan and red) to blue to white), (middle is dark, no_reflections, no_shadows), ((partial hexes on edges and sides are black))",
-    "Earth": "eight_color (tabletop_map built from small hexagon pieces) as world_map, with light_blue_is_rivers and brown_is_mountains and red_is_volcano and [white_is_snow at the top and bottom of map] as (four_color background: light_blue, green, tan, brown), horizontal_gradient is (brown to tan to green to light_blue to blue) and vertical_gradient is (white to blue to (green, tan and red) to blue to white), (middle is dark, no_reflections, no_shadows), ((partial hexes on edges and sides are black)), overhead_view, from_above",
-    "Map3": "top-down, bird's-eye view of A mystic forest-themed map viewed from directly above, showcasing detailed terrain features such as lakes, dense forests, magical flora, and hexagonal grids. The map is designed for tabletop gaming, emphasizing clarity and strategic elements, eight_color (tabletop_map built from small hexagon pieces) as alien_forest_world_map, with light_blue_is_lakes and green_is_forests and purple_is_alien_flora and [white_is_snow at the top and bottom of map] as (four_color background: light_blue, green, purple, brown), horizontal_gradient is (green to purple to light_blue to blue) and vertical_gradient is (white to blue to (green, purple and brown) to blue to white), (middle is dark, no_reflections, no_shadows), ((partial hexes on edges and sides are black))",
-    "Map4": "A top-down, bird's-eye view of a medieval battlefield map for tabletop games, eight_color (tabletop_map built from small hexagon pieces) as mystic_forest_world_map, with teal_is_lakes and dark_green_is_forest and violet_is_magic_fauna and [white_is_snow at the top and bottom of map] as (four_color background: teal, dark_green, violet, brown), horizontal_gradient is (brown to violet to teal to blue) and vertical_gradient is (white to blue to (dark_green, violet and brown) to blue to white), (middle is enchanted, no_reflections, no_shadows), ((partial hexes on edges and sides are black)), viewed from above, flat and illustrative, not a map of the Earth",
-    "Alien Landscape": "top-down view of a barren alien world map made from hexagon pieces. Features light blue rivers, brown mountains, red volcanoes, and white snow at the top and bottom. Background colors: light blue, green, tan, brown. Partial hexes on edges are black.",
-    "Alien World": "top-down view of an alien_world map built from hexagon pieces. Includes rivers, mountains, volcanoes, and snowy areas. Uses a four-color background: light blue, green, tan, brown. Partial edge hexes are black. Overhead view.",
-    "Mystic Forest": "top-down view of a mystic forest map with lakes, dense forests, magical flora, and hexagonal grids. Designed for clarity in tabletop gaming. Background colors: light blue, green, purple, brown. Partial hexes on edges are black.",
-    "Medieval Battlefield": "top-down view of a medieval battlefield map featuring lakes, forests, and magic fauna. Emphasizes clarity and strategy for tabletop games. Background colors: teal, dark green, violet, brown. Partial edge hexes are black. Viewed from above.",
+    "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.",
+    "Earth": "Top-down view of a hexagonal world map with rivers, mountains, volcanoes, and snow at top and bottom. Colors: light blue, green, tan, brown. No reflections or shadows. Partial edge hexes are black. Overhead view.",
+    "Map3": "Top-down view of a mystic forest map with lakes, dense forests, magical flora, and hex grids. Designed for tabletop gaming with clarity and strategic elements. Colors: light blue, green, purple, brown. Partial hexes on edges are black.",
+    "Map4": "Top-down view of a medieval battlefield map with lakes, forests, magical fauna, and hex grids. Emphasizes clarity and strategy for tabletop games. Colors: teal, dark green, violet, brown. Partial edge hexes are black. Viewed from above.",
+    "Alien Landscape": "Top-down view of a barren alien world map made from hexagon pieces. Features light blue rivers, brown mountains, red volcanoes, and white snow at top and bottom. Colors: light blue, green, tan, brown. Partial hexes on edges are black.",
+    "Alien World": "Top-down view of an alien world map built from hexagon pieces. Includes rivers, mountains, volcanoes, and snowy areas. Colors: light blue, green, tan, brown. Partial edge hexes are black. Overhead view.",
+    "Mystic Forest": "Top-down view of a mystic forest map with lakes, dense forests, magical flora, and hex grids. Designed for clarity in tabletop gaming. Colors: light blue, green, purple, brown. Partial hexes on edges are black.",
+    "Medieval Battlefield": "Top-down view of a medieval battlefield map featuring lakes, forests, and magic fauna. Emphasizes clarity and strategy for tabletop games. Colors: teal, dark green, violet, brown. Partial edge hexes are black. Viewed from above.",
     "Prompt": None  # Indicates that the prompt should be taken from prompt_textbox
 }
 
@@ -152,7 +158,8 @@ LORA_WEIGHTS = [
     "Cossale/Frames2-Flex.1",
     "XLabs-AI/flux-lora-collection/anime_lora.safetensors",
     "XLabs-AI/flux-lora-collection/scenery_lora.safetensors",
-    "XLabs-AI/flux-lora-collection/disney_lora.safetensors"
+    "XLabs-AI/flux-lora-collection/disney_lora.safetensors",
+    "XLabs-AI/flux-RealismLora"
 ]
 
 # Map each LoRA weight to its corresponding model
@@ -164,7 +171,8 @@ LORA_TO_MODEL = {
     "AlekseyCalvin/HSTcolorFlexAlpha": "ostris/Flex.1-alpha",
     "XLabs-AI/flux-lora-collection/anime_lora.safetensors":"black-forest-labs/FLUX.1-dev",
     "XLabs-AI/flux-lora-collection/scenery_lora.safetensors":"black-forest-labs/FLUX.1-dev",
-    "XLabs-AI/flux-lora-collection/disney_lora.safetensors":"black-forest-labs/FLUX.1-dev"
+    "XLabs-AI/flux-lora-collection/disney_lora.safetensors":"black-forest-labs/FLUX.1-dev",
+    "XLabs-AI/flux-RealismLora":"black-forest-labs/FLUX.1-dev"
 }
 condition_type = ["depth", "canny", "subject", "coloring", "deblurring", "fill", "redux"]
 # Detailed LoRA weight configurations
@@ -236,6 +244,18 @@ LORA_DETAILS = {
             }
         }
         ],
+    "XLabs-AI/flux-RealismLora":[
+        {
+            "notes": "No trigger words but 8k, Animatrix illustration style, fantasy style, natural photo cinematic should all work @6min"
+        },
+        {
+            "parameters" :{
+                "guidance_scale": "3.2",
+                "num_inference_steps": "34",
+                "scale": "0.85"
+            }
+        }
+        ],
     "Cossale/Frames2-Flex.1": [
         {
             "weight_name": "backdrops_v2.safetensors",

utils/misc.py

@@ -1,6 +1,9 @@
 # misc.py file contains miscellaneous utility functions
 import math
 import sys
+import logging
+import os
+import subprocess
 
 def pause():
     """
@@ -67,3 +70,78 @@ def convert_ratio_to_dimensions(ratio, height=512, rotate90=False):
         adjusted_width, adjusted_height = adjusted_height, adjusted_width
     return adjusted_width, adjusted_height
 
+def install_torch():    
+    print("\nInstalling PyTorch with CUDA support...")
+    # Define the package and index URL
+    package = "torch==2.4.0"
+    index_url = "https://download.pytorch.org/whl/cu124"
+    # Construct the pip install command
+    command = [
+        "pip", "install", "--force-reinstall", 
+        f"{package}", "--index-url", f"{index_url}"
+    ]
+    # Run the command using subprocess
+    subprocess.run(command, check=True)
+    print("\nPyTorch installation completed.")
+    print("\nInstalling torchvision...")
+    package = "torchvision==0.19.0"
+    index_url = "https://download.pytorch.org/whl/cu124"
+    # Construct the pip install command
+    command = [
+        "pip", "install", "--force-reinstall", 
+        f"{package}", "--index-url", f"{index_url}"
+    ]
+    # Run the command using subprocess
+    subprocess.run(command, check=True)
+    print("\nPlease restart the kernel to use the newly installed PyTorch version.")
+
+def _get_output(cmd):
+    try:
+        return subprocess.check_output(cmd).decode("utf-8")
+    except Exception as ex:
+        logging.exception(ex)
+    return None
+
+def install_cuda_toolkit():    
+    #CUDA_TOOLKIT_URL = "https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda_11.8.0_520.61.05_linux.run"
+    # CUDA_TOOLKIT_URL = "https://developer.download.nvidia.com/compute/cuda/12.2.0/local_installers/cuda_12.2.0_535.54.03_linux.run"
+    CUDA_TOOLKIT_URL = "https://developer.download.nvidia.com/compute/cuda/12.4.1/local_installers/cuda_12.4.1_550.54.15_linux.run"
+    CUDA_TOOLKIT_FILE = "/tmp/%s" % os.path.basename(CUDA_TOOLKIT_URL)
+    print("\nDownloading CUDA Toolkit from %s" % CUDA_TOOLKIT_URL)
+    subprocess.call(["wget", "-q", CUDA_TOOLKIT_URL, "-O", CUDA_TOOLKIT_FILE])
+    subprocess.call(["chmod", "+x", CUDA_TOOLKIT_FILE])
+    subprocess.call([CUDA_TOOLKIT_FILE, "--silent", "--toolkit"])    
+    os.environ["CUDA_HOME"] = "/usr/local/cuda"
+    os.environ["PATH"] = "%s/bin:%s" % (os.environ["CUDA_HOME"], os.environ["PATH"])
+    os.environ["LD_LIBRARY_PATH"] = "%s/lib:%s" % (
+        os.environ["CUDA_HOME"],
+        "" if "LD_LIBRARY_PATH" not in os.environ else os.environ["LD_LIBRARY_PATH"],
+    )
+    # Fix: arch_list[-1] += '+PTX'; IndexError: list index out of range
+    os.environ["TORCH_CUDA_ARCH_LIST"] = "8.0;8.6"
+    print("\nPlease restart the kernel to use the newly installed CUDA Toolkit.")
+
+def setup_runtime_env():
+    from torch import cuda
+    logging.info("Python Version: %s" % _get_output(["python", "--version"]))
+    logging.info("CUDA Version: %s" % _get_output(["nvcc", "--version"]))
+    logging.info("GCC Version: %s" % _get_output(["gcc", "--version"]))
+    logging.info("CUDA is available: %s" % cuda.is_available())
+    logging.info("CUDA Device Capability: %s" % (cuda.get_device_capability(),))
+
+    # Install Pre-compiled CUDA extensions (Fallback to this solution on 12/31/24)
+    # Ref: https://huggingface.co/spaces/zero-gpu-explorers/README/discussions/110
+    ##ext_dir = os.path.join(os.path.dirname(__file__), "wheels")
+    ##for e in os.listdir(ext_dir):
+    ##    logging.info("Installing Extensions from %s" % e)
+    ##    subprocess.call(
+    ##        ["pip", "install", os.path.join(ext_dir, e)], stderr=subprocess.STDOUT
+    ##    )
+    # Compile CUDA extensions
+    # Update on 12/31/24: No module named 'torch'. But it is installed and listed by `pip list`
+    # ext_dir = os.path.join(os.path.dirname(__file__), "citydreamer", "extensions")
+    # for e in os.listdir(ext_dir):
+    #     if os.path.isdir(os.path.join(ext_dir, e)):
+    #         subprocess.call(["pip", "install", "."], cwd=os.path.join(ext_dir, e))
+
+    #logging.info("Installed Python Packages: %s" % _get_output(["pip", "list"]))

utils/version_info.py

@@ -2,9 +2,8 @@
 
 import subprocess
 import os
-import spaces
-import torch
 import sys
+import gc
 import gradio as gr
 
 git = os.environ.get('GIT', "git")
@@ -46,14 +45,46 @@ def get_diffusers_version():
         return diffusers.__version__
     except Exception:
         return "<none>"
-#@spaces.GPU()
+
 def get_torch_info():
+    from torch import __version__ as torch_version_, version, cuda, backends
+    initialize_cuda()
     try:
-        return [torch.__version__, f"CUDA Version:{torch.version.cuda}", f"Available:{torch.cuda.is_available()}", f"flash attention enabled: {torch.backends.cuda.flash_sdp_enabled()}", f"Capabilities: {torch.cuda.get_device_capability(0)}", f"Device Name: {torch.cuda.get_device_name(0)}", f"Device Count: {torch.cuda.device_count()}"]
+        info = [torch_version_, f"CUDA Version:{version.cuda}", f"Available:{cuda.is_available()}", f"flash attention enabled: {backends.cuda.flash_sdp_enabled()}", f"Capabilities: {cuda.get_device_capability(0)}", f"Device Name: {cuda.get_device_name(0)}", f"Device Count: {cuda.device_count()}",f"Devices: {os.environ['CUDA_VISIBLE_DEVICES']}", f"Zero :{os.environ['CUDA_MODULE_LOADING']}"]
+        del torch_version_, version, cuda, backends
+        return info
     except Exception:
+        del torch_version_, version, cuda, backends
         return "<none>"
 
+def release_torch_resources():
+    from torch import cuda
+    # Clear the CUDA cache
+    cuda.empty_cache()
+    cuda.ipc_collect()
+    # Delete any objects that are using GPU memory
+    #for obj in gc.get_objects():
+    #    if is_tensor(obj) or (hasattr(obj, 'data') and is_tensor(obj.data)):
+    #        del obj
+    # Run garbage collection
+    del cuda
+    gc.collect()
+    
+
+def initialize_cuda():
+    from torch import cuda, version
+    if cuda.is_available():
+        device = cuda.device("cuda")
+        print(f"CUDA is available. Using device: {cuda.get_device_name(0)} with CUDA version: {version.cuda}")
+        result = "cuda"
+    else:
+        device = cuda.device("cpu")
+        print("CUDA is not available. Using CPU.")
+        result = "cpu"
+    return result
+
 def versions_html():
+    from torch import __version__ as torch_version_
     python_version = ".".join([str(x) for x in sys.version_info[0:3]])
     commit = commit_hash()
 
@@ -64,12 +95,12 @@ def versions_html():
         </a>
     '''
 
-    return f"""
+    v_html = f"""
         version: <a href="https://huggingface.co/spaces/Surn/HexaGrid/commit/{"huggingface" if commit == "<none>" else commit}" target="_blank">{"huggingface" if commit == "<none>" else commit}</a>
         &#x2000;•&#x2000;
         python: <span title="{sys.version}">{python_version}</span>
         &#x2000;•&#x2000;
-        torch: {getattr(torch, '__long_version__',torch.__version__)}
+        torch: {torch_version_}
         &#x2000;•&#x2000;
         diffusers: {get_diffusers_version()}
         &#x2000;•&#x2000;
@@ -82,4 +113,6 @@ def versions_html():
         {toggle_dark_link}
         <br>
         Full GPU Info:{get_torch_info()}
-        """
+        """
+    del torch_version_
+    return v_html