Add progress and change tokenization

app.py

@@ -6,8 +6,6 @@ from tempfile import NamedTemporaryFile
 from pathlib import Path
 import atexit
 import random
-import spaces
-
 # Import constants
 import utils.constants as constants
 
@@ -161,8 +159,7 @@ def get_model_and_lora(model_textbox):
         default_model = model_textbox
         return default_model, []
 
-#@spaces.GPU(duration=256)
-def generate_input_image_click(map_option, prompt_textbox_value, negative_prompt_textbox_value, model_textbox_value, use_conditioned_image=False, strength=0.5, image_format="16:9", scale_factor=3):
+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)):
     # Get the model and LoRA weights
     model, lora_weights = get_model_and_lora(model_textbox_value)
     global current_prerendered_image
@@ -191,7 +188,8 @@ def generate_input_image_click(map_option, prompt_textbox_value, negative_prompt
         conditioned_image,
         stength=strength,
         height=height,
-        width=width
+        width=width,
+        seed=seed
     )
        
     # Open the generated image
@@ -413,13 +411,24 @@ with gr.Blocks(css_paths="style_20250128.css", title="HexaGrid Creator", theme='
                                 label="Map Options",
                                 choices=list(constants.PROMPTS.keys()),
                                 value="Alien Landscape",
-                                elem_classes="solid"
+                                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",interactive=True)
+                            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 = 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,
@@ -571,7 +580,7 @@ with gr.Blocks(css_paths="style_20250128.css", title="HexaGrid Creator", theme='
     )
     generate_input_image.click(
         fn=generate_input_image_click,
-        inputs=[map_options, prompt_textbox, negative_prompt_textbox, model_textbox, gr.State(False), gr.State(0.5), image_size_ratio],
+        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(

utils/ai_generator.py

@@ -1,9 +1,8 @@
 # utils/ai_generator.py
-
+import gradio as gr
 import os
 import time
 from turtle import width  # Added for implementing delays
-import spaces
 import torch
 import random
 from utils.ai_generator_diffusers_flux import generate_ai_image_local
@@ -15,8 +14,7 @@ from PIL import Image
 from tempfile import NamedTemporaryFile
 import utils.constants as constants
 
-
-def generate_image_from_text(text, model_name="flax-community/dalle-mini", image_width=768, image_height=512):
+def generate_image_from_text(text, model_name="flax-community/dalle-mini", image_width=768, image_height=512, progress=gr.Progress(track_tqdm=True)):
     # Initialize the InferenceClient
     client = InferenceClient()
     # Generate the image from the text
@@ -40,12 +38,13 @@ def generate_ai_image(
     width=912,
     height=512,
     strength=0.5,
+    seed = random.randint(0, constants.MAX_SEED),
+    progress=gr.Progress(track_tqdm=True),
     *args,
     **kwargs
-):
-    seed = random.randint(1, 99999)
-    if torch.cuda.is_available():
-        print("Local GPU available. Generating image locally.")        
+):    
+    if (torch.cuda.is_available() and torch.cuda.device_count() >= 1):
+        print("Local GPU available. Generating image locally.")
         if conditioned_image is not None:
             pipeline = "FluxImg2ImgPipeline"
         return generate_ai_image_local(
@@ -69,10 +68,11 @@ def generate_ai_image(
             neg_prompt_textbox_value,
             model,
             height=height,
-            width=width
+            width=width,
+            seed=seed
         )
 
-def generate_ai_image_remote(map_option, prompt_textbox_value, neg_prompt_textbox_value, model, height=512, width=912, num_inference_steps=30, guidance_scale=3.5, seed=777):
+def generate_ai_image_remote(map_option, prompt_textbox_value, neg_prompt_textbox_value, model, height=512, width=912, num_inference_steps=30, guidance_scale=3.5, seed=777,progress=gr.Progress(track_tqdm=True)):
     max_retries = 3
     retry_delay = 4  # Initial delay in seconds
 

utils/ai_generator_diffusers_flux.py

@@ -1,11 +1,12 @@
 # utils/ai_generator_diffusers_flux.py
+import gradio as gr
 import os
 import utils.constants as constants
 import spaces
 import torch
-from diffusers import FluxPipeline,FluxImg2ImgPipeline,FluxControlPipeline, DiffusionPipeline
+from diffusers import FluxPipeline,FluxImg2ImgPipeline,FluxControlPipeline
 import accelerate 
-import transformers
+from transformers import AutoTokenizer
 import safetensors
 import xformers
 from diffusers.utils import load_image
@@ -27,7 +28,6 @@ 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
@@ -36,7 +36,6 @@ PIPELINE_CLASSES = {
     "FluxPipeline": FluxPipeline,
     "FluxImg2ImgPipeline": FluxImg2ImgPipeline
 }
-
 @spaces.GPU(duration=140)
 def generate_image_from_text(
     text,
@@ -48,16 +47,28 @@ def generate_image_from_text(
     guidance_scale=3.5,
     num_inference_steps=50,
     seed=0,
-    additional_parameters=None
+    additional_parameters=None,
+    progress=gr.Progress(track_tqdm=True)
 ):
     device = "cuda" if torch.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
     ).to(device)
-    pipe = pipe.to(device)
     pipe.enable_model_cpu_offload()
+
+    # Access the tokenizer from the pipeline
+    tokenizer = pipe.tokenizer
+
+    # Handle add_prefix_space attribute
+    if getattr(tokenizer, 'add_prefix_space', False):
+        tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)
+        # Update the pipeline's tokenizer
+        pipe.tokenizer = tokenizer
+
     # Load and apply LoRA weights
     if lora_weights:
         for lora_weight in lora_weights:
@@ -74,12 +85,18 @@ def generate_image_from_text(
                     )
             else:
                 pipe.load_lora_weights(lora_weight, use_auth_token=constants.HF_API_TOKEN)
+
+    # Set the random seed for reproducibility
     generator = torch.Generator(device=device).manual_seed(seed)
     conditions = []
+
+    # Handle conditioned image if provided
     if conditioned_image is not None:
         conditioned_image = crop_and_resize_image(conditioned_image, 1024, 1024)
         condition = Condition("subject", conditioned_image)
         conditions.append(condition)
+
+    # Prepare parameters for image generation
     generate_params = {
         "prompt": text,
         "height": image_height,
@@ -89,12 +106,24 @@ def generate_image_from_text(
         "generator": generator,
         "conditions": conditions if conditions else None
     }
+
     if additional_parameters:
         generate_params.update(additional_parameters)
     generate_params = {k: v for k, v in generate_params.items() if v is not None}
+
+    # Generate the image
     result = pipe(**generate_params)
     image = result.images[0]
     pipe.unload_lora_weights()
+
+    # Clean up
+    del result
+    del conditions
+    del generator
+    del pipe
+    torch.cuda.empty_cache()
+    torch.cuda.ipc_collect()
+
     return image
 
 @spaces.GPU(duration=140)
@@ -111,18 +140,19 @@ def generate_image_lowmem(
     seed=0,
     true_cfg_scale=1.0,
     pipeline_name="FluxPipeline",
-    strength=0.75,
-    additional_parameters=None
-):
-    print(f"\n {get_torch_info()}\n")
+    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"
     print(f"device:{device}\nmodel_name:{model_name}\nlora_weights:{lora_weights}\n")
-    
+    print(f"\n {get_torch_info()}\n")
     # Disable gradient calculations
     with torch.no_grad():
         # Initialize the pipeline inside the context manager
@@ -134,6 +164,16 @@ def generate_image_lowmem(
         pipe.enable_model_cpu_offload()
         # alternative version that may be more efficient
         # pipe.enable_sequential_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)
+            # Update the pipeline's tokenizer
+            pipe.tokenizer = tokenizer
+
         flash_attention_enabled = torch.backends.cuda.flash_sdp_enabled()
         if flash_attention_enabled == False:
             #Enable xFormers memory-efficient attention (optional)
@@ -282,6 +322,7 @@ def generate_ai_image_local (
     seed=777,
     pipeline_name="FluxPipeline",
     strength=0.75,
+    progress=gr.Progress(track_tqdm=True)
 ):
     try:
         if map_option != "Prompt":
@@ -306,10 +347,10 @@ def generate_ai_image_local (
                 additional_parameters[key] = int(value)
             elif key in ['guidance_scale','true_cfg_scale']:
                 additional_parameters[key] = float(value)
-        height = additional_parameters.get('height', height)
-        width = additional_parameters.get('width', width)
-        num_inference_steps = additional_parameters.get('num_inference_steps', num_inference_steps)
-        guidance_scale = additional_parameters.get('guidance_scale', guidance_scale)
+        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}")
@@ -347,7 +388,6 @@ def generate_ai_image_local (
         return None
 
 # does not work
-#@spaces.GPU(duration=256)
 def merge_LoRA_weights(model="black-forest-labs/FLUX.1-dev",
     lora_weights="Borcherding/FLUX.1-dev-LoRA-FractalLand-v0.1"):
     

utils/constants.py

@@ -4,6 +4,7 @@
 import os
 from pathlib import Path
 from dotenv import load_dotenv
+import numpy as np
 
 #Set the environment variables
 os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:256,expandable_segments:True"
@@ -32,6 +33,7 @@ if not HF_API_TOKEN:
     raise ValueError("HF_TOKEN is not set. Please check your .env file.")
 
 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))",

utils/live_preview_helpers.py

@@ -0,0 +1,166 @@
+import torch
+import numpy as np
+from diffusers import FluxPipeline, AutoencoderTiny, FlowMatchEulerDiscreteScheduler
+from typing import Any, Dict, List, Optional, Union
+
+# Helper functions
+def calculate_shift(
+    image_seq_len,
+    base_seq_len: int = 256,
+    max_seq_len: int = 4096,
+    base_shift: float = 0.5,
+    max_shift: float = 1.16,
+):
+    m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
+    b = base_shift - m * base_seq_len
+    mu = image_seq_len * m + b
+    return mu
+
+def retrieve_timesteps(
+    scheduler,
+    num_inference_steps: Optional[int] = None,
+    device: Optional[Union[str, torch.device]] = None,
+    timesteps: Optional[List[int]] = None,
+    sigmas: Optional[List[float]] = None,
+    **kwargs,
+):
+    if timesteps is not None and sigmas is not None:
+        raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
+    if timesteps is not None:
+        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    elif sigmas is not None:
+        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    else:
+        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+    return timesteps, num_inference_steps
+
+# FLUX pipeline function
+@torch.inference_mode()
+def flux_pipe_call_that_returns_an_iterable_of_images(
+    self,
+    prompt: Union[str, List[str]] = None,
+    prompt_2: Optional[Union[str, List[str]]] = None,
+    height: Optional[int] = None,
+    width: Optional[int] = None,
+    num_inference_steps: int = 28,
+    timesteps: List[int] = None,
+    guidance_scale: float = 3.5,
+    num_images_per_prompt: Optional[int] = 1,
+    generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+    latents: Optional[torch.FloatTensor] = None,
+    prompt_embeds: Optional[torch.FloatTensor] = None,
+    pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+    output_type: Optional[str] = "pil",
+    return_dict: bool = True,
+    joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+    max_sequence_length: int = 512,
+    good_vae: Optional[Any] = None,
+):
+    height = height or self.default_sample_size * self.vae_scale_factor
+    width = width or self.default_sample_size * self.vae_scale_factor
+
+    # 1. Check inputs
+    self.check_inputs(
+        prompt,
+        prompt_2,
+        height,
+        width,
+        prompt_embeds=prompt_embeds,
+        pooled_prompt_embeds=pooled_prompt_embeds,
+        max_sequence_length=max_sequence_length,
+    )
+
+    self._guidance_scale = guidance_scale
+    self._joint_attention_kwargs = joint_attention_kwargs
+    self._interrupt = False
+
+    # 2. Define call parameters
+    batch_size = 1 if isinstance(prompt, str) else len(prompt)
+    device = self._execution_device
+
+    # 3. Encode prompt
+    lora_scale = joint_attention_kwargs.get("scale", None) if joint_attention_kwargs is not None else None
+    prompt_embeds, pooled_prompt_embeds, text_ids = self.encode_prompt(
+        prompt=prompt,
+        prompt_2=prompt_2,
+        prompt_embeds=prompt_embeds,
+        pooled_prompt_embeds=pooled_prompt_embeds,
+        device=device,
+        num_images_per_prompt=num_images_per_prompt,
+        max_sequence_length=max_sequence_length,
+        lora_scale=lora_scale,
+    )
+    # 4. Prepare latent variables
+    num_channels_latents = self.transformer.config.in_channels // 4
+    latents, latent_image_ids = self.prepare_latents(
+        batch_size * num_images_per_prompt,
+        num_channels_latents,
+        height,
+        width,
+        prompt_embeds.dtype,
+        device,
+        generator,
+        latents,
+    )
+    # 5. Prepare timesteps
+    sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps)
+    image_seq_len = latents.shape[1]
+    mu = calculate_shift(
+        image_seq_len,
+        self.scheduler.config.base_image_seq_len,
+        self.scheduler.config.max_image_seq_len,
+        self.scheduler.config.base_shift,
+        self.scheduler.config.max_shift,
+    )
+    timesteps, num_inference_steps = retrieve_timesteps(
+        self.scheduler,
+        num_inference_steps,
+        device,
+        timesteps,
+        sigmas,
+        mu=mu,
+    )
+    self._num_timesteps = len(timesteps)
+
+    # Handle guidance
+    guidance = torch.full([1], guidance_scale, device=device, dtype=torch.float32).expand(latents.shape[0]) if self.transformer.config.guidance_embeds else None
+
+    # 6. Denoising loop
+    for i, t in enumerate(timesteps):
+        if self.interrupt:
+            continue
+
+        timestep = t.expand(latents.shape[0]).to(latents.dtype)
+
+        noise_pred = self.transformer(
+            hidden_states=latents,
+            timestep=timestep / 1000,
+            guidance=guidance,
+            pooled_projections=pooled_prompt_embeds,
+            encoder_hidden_states=prompt_embeds,
+            txt_ids=text_ids,
+            img_ids=latent_image_ids,
+            joint_attention_kwargs=self.joint_attention_kwargs,
+            return_dict=False,
+        )[0]
+        # Yield intermediate result
+        latents_for_image = self._unpack_latents(latents, height, width, self.vae_scale_factor)
+        latents_for_image = (latents_for_image / self.vae.config.scaling_factor) + self.vae.config.shift_factor
+        image = self.vae.decode(latents_for_image, return_dict=False)[0]
+        yield self.image_processor.postprocess(image, output_type=output_type)[0]
+        
+        latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+        torch.cuda.empty_cache()
+
+    # Final image using good_vae
+    latents = self._unpack_latents(latents, height, width, self.vae_scale_factor)
+    latents = (latents / good_vae.config.scaling_factor) + good_vae.config.shift_factor
+    image = good_vae.decode(latents, return_dict=False)[0]
+    self.maybe_free_model_hooks()
+    torch.cuda.empty_cache()
+    yield self.image_processor.postprocess(image, output_type=output_type)[0]