import os import time import random import logging from gradio.blocks import postprocess_update_dict import numpy as np from typing import Any, Dict, List, Optional, Union, Tuple, Iterator import spaces import gc import torch from PIL import Image import gradio as gr from tempfile import NamedTemporaryFile from diffusers import ( DiffusionPipeline, AutoencoderTiny, AutoencoderKL, AutoPipelineForImage2Image, FluxPipeline, Flux2Pipeline, AutoModel, FlowMatchEulerDiscreteScheduler, DPMSolverMultistepScheduler) from transformers import Mistral3ForConditionalGeneration from huggingface_hub import ( hf_hub_download, HfFileSystem, ModelCard, snapshot_download) from diffusers.utils import load_image from modules.version_info import ( versions_html, #initialize_cuda, #release_torch_resources, #get_torch_info ) from modules.image_utils import ( change_color, open_image, build_prerendered_images_by_quality, upscale_image, # lerp_imagemath, # shrink_and_paste_on_blank, show_lut, apply_lut_to_image_path, multiply_and_blend_images, alpha_composite_with_control, apply_alpha_mask, resize_and_crop_image, convert_to_rgba_png, get_image_from_dict ) from modules.constants import ( LORA_DETAILS, LORAS as loras, MODELS, LORA_TO_MODEL, default_lut_example_img, lut_files, MAX_SEED, IS_SHARED_SPACE, # lut_folder,cards, # cards_alternating, # card_colors, # card_colors_alternating, pre_rendered_maps_paths, PROMPTS, NEGATIVE_PROMPTS, TARGET_SIZE, temp_files, load_env_vars, dotenv_path ) from modules.color_utils import ( color_to_hex ) # from modules.excluded_colors import ( # add_color, # delete_color, # build_dataframe, # on_input, # excluded_color_list, # on_color_display_select # ) from modules.misc import ( get_filename, convert_ratio_to_dimensions, update_dimensions_on_ratio, wait5 ) from modules.lora_details import ( approximate_token_count, split_prompt_precisely, upd_prompt_notes_by_index, get_trigger_words_by_index ) from modules.mazlib.cli import run_maze as run_maze_maker input_image_palette = [] current_prerendered_image = gr.State("./images/Beeuty-1.png") user_info = { "username": "guest", "session_hash": None, "headers": None, "client": None, "query_params": None, "path_params": None, "level" : 0 } user_profile = None def get_profile() -> gr.OAuthProfile | None: global user_profile return user_profile # Define a function to handle the login button click and retrieve user information. def handle_login(profile: gr.OAuthProfile | None, request: gr.Request): # Extract user information from the request if profile is None: profile = gr.State(None) profile = get_profile() user_info = { "username": request.username, "session_hash": request.session_hash, "headers": dict(request.headers), "client": request.client, "query_params": dict(request.query_params), "path_params": dict(request.path_params), "level" : (0 if request.username == "guest" else 2) } return user_info, gr.update(logout_value=f"Logout {user_info['username']} ({user_info['level']})", value=f"Login {user_info['username']} ({user_info['level']})"), profile #---if workspace = local or colab--- # Authenticate with Hugging Face # from huggingface_hub import login # Log in to Hugging Face using the provided token # hf_token = 'hf-token-authentication' # login(hf_token) 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 def is_flux2_model(model_name: str) -> bool: return "FLUX.2" in model_name.upper() or model_name.lower().startswith("fal/flux.2") def get_total_vram_gb(device_index: int = 0) -> float: if not torch.cuda.is_available(): return 0.0 try: total_bytes = torch.cuda.get_device_properties(device_index).total_memory return float(total_bytes) / (1024 ** 3) except Exception: return 0.0 def should_enable_flux2_cpu_offload(vram_threshold_gb: float = 80.0) -> bool: return get_total_vram_gb() < vram_threshold_gb def _coerce_lora_parameter_value(value): if isinstance(value, str): v = value.strip() try: if "." in v: return float(v) return int(v) except ValueError: return value return value def get_lora_runtime_parameters(lora_repo: str, selected_lora: Optional[Dict[str, Any]] = None) -> Dict[str, Any]: params: Dict[str, Any] = {} detail_items = LORA_DETAILS.get(lora_repo, []) for item in detail_items: if isinstance(item, dict): detail_params = item.get("parameters") if isinstance(detail_params, dict): for key, value in detail_params.items(): params[key] = _coerce_lora_parameter_value(value) if selected_lora and isinstance(selected_lora.get("parameters"), dict): for key, value in selected_lora["parameters"].items(): params[key] = _coerce_lora_parameter_value(value) return params # Preload settings FLUX2_PRELOAD_REPO = "diffusers/FLUX.2-dev-bnb-4bit" FLUX2_PRELOAD_PATTERNS = [ "model_index.json", "transformer/*", "vae/*", "scheduler/*", "tokenizer/*", "tokenizer_2/*", "text_encoder/*", "text_encoder_2/*", ] FLUX2_TURBO_LORA_REPO = "fal/FLUX.2-dev-Turbo" FLUX2_TURBO_LORA_FILE = "flux.2-turbo-lora.safetensors" def preload_flux2_model_to_cache() -> None: """Warm HF cache at startup (no model instantiation).""" token = os.getenv("HF_TOKEN", None) # 1) Base FLUX.2 model cache warmup t0 = time.time() flux2_repo_id = FLUX2_PRELOAD_REPO.strip() print(f"[Startup] Preloading {flux2_repo_id} into Hugging Face cache...") try: local_path = snapshot_download( repo_id=flux2_repo_id, allow_patterns=FLUX2_PRELOAD_PATTERNS, resume_download=True, token=token, ) print(f"[Startup] FLUX.2 cache preload complete in {time.time() - t0:.2f}s") print(f"[Startup] FLUX.2 cached at: {local_path}") except Exception as ex: print(f"[Startup] FLUX.2 preload skipped/failed after {time.time() - t0:.2f}s: {ex}") # 2) FLUX.2 Turbo LoRA cache warmup t1 = time.time() print(f"[Startup] Preloading {FLUX2_TURBO_LORA_REPO}/{FLUX2_TURBO_LORA_FILE}...") try: lora_path = hf_hub_download( repo_id=FLUX2_TURBO_LORA_REPO, filename=FLUX2_TURBO_LORA_FILE, token=token, ) print(f"[Startup] FLUX.2 Turbo LoRA cached in {time.time() - t1:.2f}s") print(f"[Startup] FLUX.2 Turbo LoRA path: {lora_path}") except Exception as ex: print(f"[Startup] FLUX.2 Turbo LoRA preload skipped/failed after {time.time() - t1:.2f}s: {ex}") @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, negative_prompt: 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, ): # Set default height and width height = height or self.default_sample_size * self.vae_scale_factor width = width or self.default_sample_size * self.vae_scale_factor # Validate 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 batch_size = 1 if isinstance(prompt, str) else len(prompt) device = self._execution_device lora_scale = joint_attention_kwargs.get("scale", None) if joint_attention_kwargs is not None else None # Prepare prompt inputs. Flux2Pipeline does not accept prompt_2 or negative_prompt # as separate kwargs, so merge them into a single prompt string for Flux2. if isinstance(self, Flux2Pipeline): combined_prompt = None if prompt is None: combined_prompt = prompt_2 elif prompt_2 is None: combined_prompt = prompt else: combined_prompt = f"{prompt} {prompt_2}" prompt_to_encode = combined_prompt neg_prompt_to_encode = None else: prompt_to_encode = prompt neg_prompt_to_encode = None if negative_prompt is None else negative_prompt # Encode the positive prompt prompt_embeds_pos, pooled_prompt_embeds_pos, text_ids_pos = self.encode_prompt( prompt=prompt_to_encode, prompt_2=None, 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, ) # Encode the negative prompt if provided # Encode negative prompt only for pipelines that support it if (not isinstance(self, Flux2Pipeline)) and (negative_prompt is not None): prompt_embeds_neg, pooled_prompt_embeds_neg, text_ids_neg = self.encode_prompt( prompt=negative_prompt, prompt_2=None, prompt_embeds=None, pooled_prompt_embeds=None, device=device, num_images_per_prompt=num_images_per_prompt, max_sequence_length=max_sequence_length, lora_scale=lora_scale, ) else: # Fallback to positive embeddings if no negative prompt or unsupported prompt_embeds_neg = prompt_embeds_pos pooled_prompt_embeds_neg = pooled_prompt_embeds_pos text_ids_neg = text_ids_pos # Prepare latents 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_pos.dtype, device, generator, latents, ) # Set up 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) guidance = ( torch.full([1], guidance_scale, device=device, dtype=torch.float32).expand(latents.shape[0]) if self.transformer.config.guidance_embeds else None ) # Denoising loop for i, t in enumerate(timesteps): if self._interrupt: continue timestep = t.expand(latents.shape[0]).to(latents.dtype) print(f"Step {i + 1}/{num_inference_steps} - Timestep: {timestep.item()}\n") # Compute noise prediction for positive prompt noise_pred_pos = self.transformer( hidden_states=latents, timestep=timestep / 1000, guidance=guidance, pooled_projections=pooled_prompt_embeds_pos, encoder_hidden_states=prompt_embeds_pos, txt_ids=text_ids_pos, img_ids=latent_image_ids, joint_attention_kwargs=self.joint_attention_kwargs, return_dict=False, )[0] # Compute noise prediction for negative prompt noise_pred_neg = self.transformer( hidden_states=latents, timestep=timestep / 1000, guidance=guidance, pooled_projections=pooled_prompt_embeds_neg, encoder_hidden_states=prompt_embeds_neg, txt_ids=text_ids_neg, img_ids=latent_image_ids, joint_attention_kwargs=self.joint_attention_kwargs, return_dict=False, )[0] # Combine noise predictions using guidance scale noise_pred = noise_pred_neg + guidance_scale * (noise_pred_pos - noise_pred_neg) # Generate intermediate image 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] # Update latents with combined noise prediction latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0] torch.cuda.empty_cache() # Final image generation 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] @torch.inference_mode() def flux_pipe_call_single_final_image( self, prompt: Union[str, List[str]] = None, prompt_2: Optional[Union[str, List[str]]] = None, negative_prompt: Optional[Union[str, List[str]]] = None, height: Optional[int] = None, width: Optional[int] = None, num_inference_steps: int = 28, guidance_scale: float = 3.5, generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, joint_attention_kwargs: Optional[Dict[str, Any]] = None, output_type: Optional[str] = "pil", ): is_f2 = isinstance(self, Flux2Pipeline) final_steps = num_inference_steps if num_inference_steps is not None else (7 if is_f2 else 28) kwargs = { "prompt": prompt, "height": height, "width": width, "num_inference_steps": final_steps, "guidance_scale": guidance_scale, "generator": generator, "output_type": output_type, } if is_f2: # Flux2Pipeline does not accept prompt_2; merge into prompt. if prompt is None: kwargs["prompt"] = prompt_2 elif prompt_2 is not None: kwargs["prompt"] = f"{prompt} {prompt_2}" else: if prompt_2 is not None: kwargs["prompt_2"] = prompt_2 # Flux2Pipeline does not support negative_prompt if (negative_prompt is not None) and (not is_f2): kwargs["negative_prompt"] = negative_prompt if joint_attention_kwargs is not None: if is_f2: kwargs["attention_kwargs"] = joint_attention_kwargs else: kwargs["joint_attention_kwargs"] = joint_attention_kwargs return self(**kwargs).images[0] #--------------------------------------------------Model Initialization-----------------------------------------------------------------------------------------# # Pre-download FLUX.2 once at startup so later switching is cache-hit. preload_flux2_model_to_cache() dtype = torch.bfloat16 device = "cpu" # ZeroGPU-safe at startup; move to CUDA only inside @spaces.GPU base_model = "black-forest-labs/FLUX.1-dev" active_base_model = base_model flux2_cpu_offload_enabled = False taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype) good_vae = AutoencoderKL.from_pretrained(base_model, subfolder="vae", torch_dtype=dtype) pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=dtype, vae=taef1) pipe_i2i = None # pipe_i2i = AutoPipelineForImage2Image.from_pretrained(base_model, # vae=good_vae, # transformer=pipe.transformer, # text_encoder=pipe.text_encoder, # tokenizer=pipe.tokenizer, # text_encoder_2=pipe.text_encoder_2, # tokenizer_2=pipe.tokenizer_2, # torch_dtype=dtype # ) pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe) pipe.flux_pipe_call_single_final_image = flux_pipe_call_single_final_image.__get__(pipe) def _release_gpu_obj(obj): if obj is None: return None try: if hasattr(obj, "to"): obj.to("cpu") except Exception: pass del obj gc.collect() # IMPORTANT: no torch.cuda.* calls here (can run outside ZeroGPU allowance) return None def ensure_pipelines_for_model(model_name: str, use_safetensors: bool = False, image_input=None): global pipe, pipe_i2i, good_vae, active_base_model, flux2_cpu_offload_enabled need_i2i = image_input is not None already_correct_model = (model_name == active_base_model) # Fast path: correct model and requested variant already exists. if already_correct_model: if need_i2i and pipe_i2i is not None: pipe = _release_gpu_obj(pipe) # enforce exclusivity status_msg = f"Using base model: {active_base_model} (i2i ready)" print(status_msg) gr.Info(status_msg) return if (not need_i2i) and pipe is not None: pipe_i2i = _release_gpu_obj(pipe_i2i) # enforce exclusivity status_msg = f"Using base model: {active_base_model} (t2i ready)" print(status_msg) gr.Info(status_msg) return status_msg = f"Rebuilding {'i2i' if need_i2i else 't2i'} pipeline for {model_name}" print(status_msg) gr.Info(status_msg) else: status_msg = f"Switching base model from {active_base_model} to {model_name}" print(status_msg) gr.Info(status_msg) # release old model resources before rebuilding pipe = _release_gpu_obj(pipe) pipe_i2i = _release_gpu_obj(pipe_i2i) good_vae = _release_gpu_obj(good_vae) # Build t2i source pipe if needed (also needed as source when creating i2i). if pipe is None: if is_flux2_model(model_name): text_encoder = Mistral3ForConditionalGeneration.from_pretrained( model_name, subfolder="text_encoder", torch_dtype=torch.bfloat16, device_map="cpu", local_files_only=True, ) dit = AutoModel.from_pretrained( model_name, subfolder="transformer", torch_dtype=torch.bfloat16, device_map="cpu", local_files_only=True, ) pipe = Flux2Pipeline.from_pretrained( model_name, torch_dtype=dtype, text_encoder=text_encoder, transformer=dit, use_safetensors=use_safetensors, local_files_only=True, ) flux2_cpu_offload_enabled = should_enable_flux2_cpu_offload(60.0) if flux2_cpu_offload_enabled: pipe.enable_model_cpu_offload() print(f"FLUX.2 CPU offload enabled (VRAM: {get_total_vram_gb():.2f} GB)") else: print(f"FLUX.2 CPU offload disabled (VRAM: {get_total_vram_gb():.2f} GB)") good_vae = pipe.vae elif "FLUX" in model_name.upper(): flux2_cpu_offload_enabled = False good_vae = AutoencoderKL.from_pretrained( model_name, subfolder="vae", torch_dtype=dtype, use_safetensors=use_safetensors ) pipe = FluxPipeline.from_pretrained( model_name, torch_dtype=dtype, vae=taef1, use_safetensors=use_safetensors ) else: flux2_cpu_offload_enabled = False good_vae = AutoencoderKL.from_pretrained( model_name, subfolder="vae", torch_dtype=dtype, use_safetensors=use_safetensors ) pipe = DiffusionPipeline.from_pretrained( model_name, torch_dtype=dtype, vae=taef1, use_safetensors=use_safetensors ) pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe) pipe.flux_pipe_call_single_final_image = flux_pipe_call_single_final_image.__get__(pipe) if need_i2i: pipe_i2i = _release_gpu_obj(pipe_i2i) if is_flux2_model(model_name): # FLUX.2 i2i: reuse Flux2Pipeline (no AutoPipelineForImage2Image) pipe_i2i = pipe pipe = None else: pipe_i2i_kwargs = { "vae": good_vae, "transformer": pipe.transformer, "torch_dtype": dtype, "use_safetensors": use_safetensors, "local_files_only": True, } if hasattr(pipe, "text_encoder") and pipe.text_encoder is not None: pipe_i2i_kwargs["text_encoder"] = pipe.text_encoder if hasattr(pipe, "tokenizer") and pipe.tokenizer is not None: pipe_i2i_kwargs["tokenizer"] = pipe.tokenizer if hasattr(pipe, "text_encoder_2") and pipe.text_encoder_2 is not None: pipe_i2i_kwargs["text_encoder_2"] = pipe.text_encoder_2 if hasattr(pipe, "tokenizer_2") and pipe.tokenizer_2 is not None: pipe_i2i_kwargs["tokenizer_2"] = pipe.tokenizer_2 pipe_i2i = AutoPipelineForImage2Image.from_pretrained(model_name, **pipe_i2i_kwargs) pipe = _release_gpu_obj(pipe) else: pipe_i2i = _release_gpu_obj(pipe_i2i) active_base_model = model_name def _move_for_inference(main_pipe, vae_obj): if isinstance(main_pipe, Flux2Pipeline) and flux2_cpu_offload_enabled: return if torch.cuda.is_available(): try: main_pipe.to("cuda") vae_obj.to("cuda") except RuntimeError as ex: if "CUDACachingAllocator" in str(ex) or "NVML_SUCCESS" in str(ex): torch.cuda.empty_cache() try: torch.cuda.ipc_collect() except Exception: pass main_pipe.to("cuda") vae_obj.to("cuda") else: raise class calculateDuration: def __init__(self, activity_name=""): self.activity_name = activity_name def __enter__(self): self.start_time = time.time() return self def __exit__(self, exc_type, exc_value, traceback): self.end_time = time.time() self.elapsed_time = self.end_time - self.start_time if self.activity_name: print(f"Elapsed time for {self.activity_name}: {self.elapsed_time:.6f} seconds") else: print(f"Elapsed time: {self.elapsed_time:.6f} seconds") def update_selection(evt: gr.SelectData, width, height, aspect_ratio, current_steps): selected_lora = loras[evt.index] new_placeholder = f"Type a prompt for {selected_lora['title']}" new_aspect_ratio = aspect_ratio lora_repo = selected_lora["repo"] target_model = LORA_TO_MODEL.get(lora_repo, base_model) lora_runtime_parameters = get_lora_runtime_parameters(lora_repo, selected_lora) lora_default_steps = lora_runtime_parameters.get("num_inference_steps") if isinstance(lora_default_steps, float): lora_default_steps = int(lora_default_steps) if is_flux2_model(target_model): flux2_steps = lora_default_steps if isinstance(lora_default_steps, int) else 7 steps_update = gr.update( minimum=flux2_steps, maximum=flux2_steps, step=1, value=flux2_steps, interactive=False, label="Steps (from LoRA defaults for FLUX.2)", ) else: safe_steps = 7 try: if current_steps is not None: safe_steps = int(current_steps) except Exception: safe_steps = 7 steps_update = gr.update( minimum=1, maximum=50, step=1, value=max(1, min(50, safe_steps)), interactive=True, label="Steps", ) updated_text = f"### Selected: [{lora_repo}](https://huggingface.co/{lora_repo}) ✅" # aspect will now use ratios if implemented, like 16:9, 4:3, 1:1 if "aspect" in selected_lora: try: new_aspect_ratio = selected_lora["aspect"] width, height = update_dimensions_on_ratio(new_aspect_ratio, height) except Exception as e: print(f"\nError in update selection aspect ratios:{e}\nSkipping") new_aspect_ratio = aspect_ratio width = width height = height return ( gr.update(placeholder=new_placeholder), updated_text, evt.index, width, height, new_aspect_ratio, upd_prompt_notes_by_index(evt.index), steps_update, ) @spaces.GPU(duration=120,progress=gr.Progress(track_tqdm=True)) def generate_image(prompt_mash, negative_prompt, steps, seed, cfg_scale, width, height, lora_scale, progress): _move_for_inference(pipe, good_vae) generator = torch.Generator(device="cuda").manual_seed(seed) flash_attention_enabled = torch.backends.cuda.flash_sdp_enabled() if flash_attention_enabled: pipe.attn_implementation="flash_attention_2" if IS_SHARED_SPACE: pipe.vae.enable_tiling() # For larger resolutions if needed else: # Compile UNet #pipe.transformer = torch.compile(pipe.transformer, mode="reduce-overhead") #pipe.enable_model_cpu_offload() #for smaller GPUs pipe.vae.enable_slicing() # Disable unnecessary features pipe.safety_checker = None print(f"\nGenerating image with prompt: {prompt_mash}\n") approx_tokens= approximate_token_count(prompt_mash) if approx_tokens > 76: print(f"\nSplitting prompt due to length: {approx_tokens}\n") prompt, prompt2 = split_prompt_precisely(prompt_mash) else: prompt = prompt_mash prompt2 = None with calculateDuration("Generating image"): # Generate image if isinstance(pipe, Flux2Pipeline): final_image = pipe.flux_pipe_call_single_final_image( prompt=prompt, prompt_2=prompt2, num_inference_steps=steps, guidance_scale=cfg_scale, width=width, height=height, generator=generator, joint_attention_kwargs={"scale": lora_scale}, output_type="pil", ) yield final_image else: for img in pipe.flux_pipe_call_that_returns_an_iterable_of_images( prompt=prompt, prompt_2=prompt2, num_inference_steps=steps, guidance_scale=cfg_scale, width=width, height=height, generator=generator, joint_attention_kwargs={"scale": lora_scale}, output_type="pil", good_vae=good_vae, ): yield img @spaces.GPU(duration=120,progress=gr.Progress(track_tqdm=True)) def generate_image_to_image(prompt_mash, negative_prompt, image_input_path, image_strength, steps, cfg_scale, width, height, lora_scale, seed, progress): _move_for_inference(pipe_i2i, good_vae) generator = torch.Generator(device="cuda").manual_seed(seed) flash_attention_enabled = torch.backends.cuda.flash_sdp_enabled() if flash_attention_enabled: pipe_i2i.attn_implementation="flash_attention_2" if IS_SHARED_SPACE: pipe_i2i.vae.enable_tiling() # For larger resolutions if needed else: # Compile UNet # pipe.transformer = torch.compile(pipe.transformer, mode="reduce-overhead") # uses the other pipe's transformer # pipe_i2i.enable_model_cpu_offload() #for smaller GPUs pipe_i2i.vae.enable_slicing() # Disable unnecessary features pipe_i2i.safety_checker = None image_input = open_image(image_input_path) print(f"\nGenerating image with prompt: {prompt_mash} and {image_input_path}\n") approx_tokens= approximate_token_count(prompt_mash) if approx_tokens > 76: print(f"\nSplitting prompt due to length: {approx_tokens}\n") prompt, prompt2 = split_prompt_precisely(prompt_mash) else: prompt = prompt_mash prompt2 = None with calculateDuration("Generating image"): # Generate image if isinstance(pipe_i2i, Flux2Pipeline): # Flux2Pipeline does not accept prompt_2; merge prompt and prompt_2 if prompt is None: merged_prompt = prompt2 elif prompt2 is None: merged_prompt = prompt else: merged_prompt = f"{prompt} {prompt2}" final_image = pipe_i2i( prompt=merged_prompt, image=image_input, num_inference_steps=steps, guidance_scale=cfg_scale, width=width, height=height, generator=generator, attention_kwargs={"scale": lora_scale}, output_type="pil", ).images[0] else: final_image = pipe_i2i( prompt=prompt, prompt_2=prompt2, image=image_input, strength=image_strength, num_inference_steps=steps, guidance_scale=cfg_scale, width=width, height=height, generator=generator, joint_attention_kwargs={"scale": lora_scale}, output_type="pil", ).images[0] return final_image def run_lora(prompt, map_option, image_input, image_strength, cfg_scale, steps, selected_index, randomize_seed, seed, width, height, lora_scale, enlarge, use_conditioned_image=False, progress=gr.Progress(track_tqdm=True)) -> Iterator[Tuple[str, int, Any]]: if selected_index is None: raise gr.Error("You must select a LoRA before proceeding.🧨") print(f"input Image: {image_input}\n") # handle selecting a conditioned image from the gallery global current_prerendered_image conditioned_image=None formatted_map_option = map_option.lower().replace(' ', '_') 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): conditioned_image = open_image(current_prerendered_image.value).convert("RGB") image_input = resize_and_crop_image(conditioned_image, width, height) print(f"Conditioned Image: {image_input.size}.. converted to RGB and resized\n") if map_option != "Prompt": prompt = PROMPTS[map_option] negative_prompt = NEGATIVE_PROMPTS.get(map_option, "") selected_lora = loras[selected_index] lora_path = selected_lora["repo"] target_model = LORA_TO_MODEL.get(lora_path, base_model) lora_runtime_parameters = get_lora_runtime_parameters(lora_path, selected_lora) # detect safetensor usage use_safetensors_flag = False # check explicit "weights" field on the lora entry weight_field = selected_lora.get("weights") or "" if weight_field and "safetensor" in weight_field.lower(): use_safetensors_flag = True else: # fallback: check LORA_DETAILS for weight_name entries details = LORA_DETAILS.get(lora_path) if details: for d in details: wn = d.get("weight_name") or d.get("weights") or "" if wn and "safetensor" in wn.lower(): use_safetensors_flag = True break ensure_pipelines_for_model(target_model, use_safetensors=use_safetensors_flag, image_input=image_input) runtime_steps = lora_runtime_parameters.get("num_inference_steps") runtime_guidance = lora_runtime_parameters.get("guidance_scale") if isinstance(runtime_steps, (int, float)): steps = int(runtime_steps) elif is_flux2_model(active_base_model): steps = 8 if isinstance(runtime_guidance, (int, float)): cfg_scale = float(runtime_guidance) trigger_word = selected_lora["trigger_word"] if(trigger_word): if "trigger_position" in selected_lora: if selected_lora["trigger_position"] == "prepend": prompt_mash = f"{trigger_word} {prompt}" else: prompt_mash = f"{prompt} {trigger_word}" else: prompt_mash = f"{trigger_word} {prompt}" else: prompt_mash = prompt with calculateDuration("Unloading LoRA"): if pipe is not None: pipe.unload_lora_weights() if pipe_i2i is not None: pipe_i2i.unload_lora_weights() #LoRA weights flow with calculateDuration(f"Loading LoRA weights for {selected_lora['title']}"): pipe_to_use = pipe_i2i if image_input is not None else pipe weight_name = selected_lora.get("weights", None) pipe_to_use.load_lora_weights( lora_path, weight_name=weight_name, low_cpu_mem_usage=True ) with calculateDuration("Randomizing seed"): if randomize_seed: seed = random.randint(0, MAX_SEED) if(image_input is not None): print(f"\nGenerating image to image with seed: {seed}\n") generated_image = generate_image_to_image(prompt_mash, negative_prompt, image_input, image_strength, steps, cfg_scale, width, height, lora_scale, seed, progress) if enlarge: upscale_factor = max(1.0, min((TARGET_SIZE[0] / width), (TARGET_SIZE[1] / height))) progress(0.75, desc=f"upscaling to {upscale_factor}") upscaled_image = upscale_image(generated_image, upscale_factor) progress(1.0, desc=f"upscaled to {upscale_factor}") else: upscaled_image = generated_image # Save the upscaled image to a temporary file with NamedTemporaryFile(delete=False, suffix=".png", prefix=f"{formatted_map_option}_") as tmp_upscaled: upscaled_image.save(tmp_upscaled.name, format="PNG") temp_files.append(tmp_upscaled.name) print(f"Upscaled image saved to {tmp_upscaled.name}") final_image = tmp_upscaled.name yield final_image, seed, gr.update(visible=False) else: image_generator = generate_image(prompt_mash, negative_prompt, steps, seed, cfg_scale, width, height, lora_scale, progress) final_image = None step_counter = 0 # create a single temporary file for intermediate previews with NamedTemporaryFile(delete=False, suffix=".png", prefix=f"{formatted_map_option}_step_") as tmp_step: tmp_preview_path = tmp_step.name for image in image_generator: step_counter += 1 generated_image = image # PIL.Image.Image # overwrite the same preview file to update UI without creating many files try: generated_image.save(tmp_preview_path, format="PNG") except Exception: # fallback: save to a new temp file if overwrite fails with NamedTemporaryFile(delete=False, suffix=".png", prefix=f"{formatted_map_option}_step_") as tf: generated_image.save(tf.name, format="PNG") tmp_preview_path = tf.name progress_bar = f'
"+trigger_word+" as the trigger word" if trigger_word else "No trigger word found. If there's a trigger word, include it in your prompt"}