Surn/HexGameMaker
4
1import os2import time3import random4import logging5from gradio.blocks import postprocess_update_dict6import numpy as np7from typing import Any, Dict, List, Optional, Union, Tuple, Iterator8 9import spaces10 11import gc12import torch13from PIL import Image14import gradio as gr15from tempfile import NamedTemporaryFile16 17from diffusers import (18 DiffusionPipeline,19 AutoencoderTiny,20 AutoencoderKL,21 AutoPipelineForImage2Image,22 FluxPipeline, Flux2Pipeline, AutoModel,23 FlowMatchEulerDiscreteScheduler,24 DPMSolverMultistepScheduler)25from transformers import Mistral3ForConditionalGeneration26 27from huggingface_hub import (28 hf_hub_download,29 HfFileSystem,30 ModelCard,31 snapshot_download)32 33from diffusers.utils import load_image34 35from modules.version_info import (36 versions_html,37 #initialize_cuda,38 #release_torch_resources,39 #get_torch_info40)41from modules.image_utils import (42 change_color,43 open_image,44 build_prerendered_images_by_quality,45 upscale_image,46 # lerp_imagemath,47 # shrink_and_paste_on_blank,48 show_lut,49 apply_lut_to_image_path,50 multiply_and_blend_images,51 alpha_composite_with_control,52 apply_alpha_mask,53 resize_and_crop_image,54 convert_to_rgba_png,55 get_image_from_dict56)57from modules.constants import (58 LORA_DETAILS, LORAS as loras, MODELS, LORA_TO_MODEL,59 default_lut_example_img, 60 lut_files, 61 MAX_SEED, 62 IS_SHARED_SPACE,63 # lut_folder,cards, 64 # cards_alternating, 65 # card_colors, 66 # card_colors_alternating,67 pre_rendered_maps_paths,68 PROMPTS,69 NEGATIVE_PROMPTS,70 TARGET_SIZE,71 temp_files,72 load_env_vars,73 dotenv_path74)75from modules.color_utils import (76 color_to_hex77)78 79# from modules.excluded_colors import (80# add_color,81# delete_color,82# build_dataframe,83# on_input,84# excluded_color_list,85# on_color_display_select86# )87from modules.misc import (88 get_filename,89 convert_ratio_to_dimensions,90 update_dimensions_on_ratio,91 wait592)93from modules.lora_details import (94 approximate_token_count,95 split_prompt_precisely,96 upd_prompt_notes_by_index,97 get_trigger_words_by_index98)99 100from modules.mazlib.cli import run_maze as run_maze_maker101 102input_image_palette = []103current_prerendered_image = gr.State("./images/Beeuty-1.png")104user_info = {105 "username": "guest",106 "session_hash": None,107 "headers": None,108 "client": None,109 "query_params": None,110 "path_params": None,111 "level" : 0112}113 114user_profile = None115 116def get_profile() -> gr.OAuthProfile | None:117 global user_profile118 return user_profile119 120# Define a function to handle the login button click and retrieve user information.121def handle_login(profile: gr.OAuthProfile | None, request: gr.Request): 122 # Extract user information from the request 123 if profile is None:124 profile = gr.State(None)125 profile = get_profile()126 user_info = {127 "username": request.username,128 "session_hash": request.session_hash,129 "headers": dict(request.headers),130 "client": request.client,131 "query_params": dict(request.query_params),132 "path_params": dict(request.path_params),133 "level" : (0 if request.username == "guest" else 2)134 }135 return user_info, gr.update(logout_value=f"Logout {user_info['username']} ({user_info['level']})", value=f"Login {user_info['username']} ({user_info['level']})"), profile136#---if workspace = local or colab---137 138# Authenticate with Hugging Face139# from huggingface_hub import login140 141# Log in to Hugging Face using the provided token142# hf_token = 'hf-token-authentication'143# login(hf_token)144 145def calculate_shift(146 image_seq_len,147 base_seq_len: int = 256,148 max_seq_len: int = 4096,149 base_shift: float = 0.5,150 max_shift: float = 1.16,151):152 m = (max_shift - base_shift) / (max_seq_len - base_seq_len)153 b = base_shift - m * base_seq_len154 mu = image_seq_len * m + b155 return mu156 157def retrieve_timesteps(158 scheduler,159 num_inference_steps: Optional[int] = None,160 device: Optional[Union[str, torch.device]] = None,161 timesteps: Optional[List[int]] = None,162 sigmas: Optional[List[float]] = None,163 **kwargs,164):165 if timesteps is not None and sigmas is not None:166 raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")167 if timesteps is not None:168 scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)169 timesteps = scheduler.timesteps170 num_inference_steps = len(timesteps)171 elif sigmas is not None:172 scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)173 timesteps = scheduler.timesteps174 num_inference_steps = len(timesteps)175 else:176 scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)177 timesteps = scheduler.timesteps178 return timesteps, num_inference_steps179 180 181def is_flux2_model(model_name: str) -> bool:182 return "FLUX.2" in model_name.upper() or model_name.lower().startswith("fal/flux.2")183 184 185def get_total_vram_gb(device_index: int = 0) -> float:186 if not torch.cuda.is_available():187 return 0.0188 try:189 total_bytes = torch.cuda.get_device_properties(device_index).total_memory190 return float(total_bytes) / (1024 ** 3)191 except Exception:192 return 0.0193 194 195def should_enable_flux2_cpu_offload(vram_threshold_gb: float = 80.0) -> bool:196 return get_total_vram_gb() < vram_threshold_gb197 198 199def _coerce_lora_parameter_value(value):200 if isinstance(value, str):201 v = value.strip()202 try:203 if "." in v:204 return float(v)205 return int(v)206 except ValueError:207 return value208 return value209 210 211def get_lora_runtime_parameters(lora_repo: str, selected_lora: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:212 params: Dict[str, Any] = {}213 214 detail_items = LORA_DETAILS.get(lora_repo, [])215 for item in detail_items:216 if isinstance(item, dict):217 detail_params = item.get("parameters")218 if isinstance(detail_params, dict):219 for key, value in detail_params.items():220 params[key] = _coerce_lora_parameter_value(value)221 222 if selected_lora and isinstance(selected_lora.get("parameters"), dict):223 for key, value in selected_lora["parameters"].items():224 params[key] = _coerce_lora_parameter_value(value)225 226 return params227 228# Preload settings229FLUX2_PRELOAD_REPO = "diffusers/FLUX.2-dev-bnb-4bit"230FLUX2_PRELOAD_PATTERNS = [231 "model_index.json",232 "transformer/*",233 "vae/*",234 "scheduler/*",235 "tokenizer/*",236 "tokenizer_2/*",237 "text_encoder/*",238 "text_encoder_2/*",239]240 241FLUX2_TURBO_LORA_REPO = "fal/FLUX.2-dev-Turbo"242FLUX2_TURBO_LORA_FILE = "flux.2-turbo-lora.safetensors"243 244 245def preload_flux2_model_to_cache() -> None:246 """Warm HF cache at startup (no model instantiation)."""247 token = os.getenv("HF_TOKEN", None)248 249 # 1) Base FLUX.2 model cache warmup250 t0 = time.time()251 flux2_repo_id = FLUX2_PRELOAD_REPO.strip()252 print(f"[Startup] Preloading {flux2_repo_id} into Hugging Face cache...")253 try:254 local_path = snapshot_download(255 repo_id=flux2_repo_id,256 allow_patterns=FLUX2_PRELOAD_PATTERNS,257 resume_download=True,258 token=token,259 )260 print(f"[Startup] FLUX.2 cache preload complete in {time.time() - t0:.2f}s")261 print(f"[Startup] FLUX.2 cached at: {local_path}")262 except Exception as ex:263 print(f"[Startup] FLUX.2 preload skipped/failed after {time.time() - t0:.2f}s: {ex}")264 265 # 2) FLUX.2 Turbo LoRA cache warmup266 t1 = time.time()267 print(f"[Startup] Preloading {FLUX2_TURBO_LORA_REPO}/{FLUX2_TURBO_LORA_FILE}...")268 try:269 lora_path = hf_hub_download(270 repo_id=FLUX2_TURBO_LORA_REPO,271 filename=FLUX2_TURBO_LORA_FILE,272 token=token,273 )274 print(f"[Startup] FLUX.2 Turbo LoRA cached in {time.time() - t1:.2f}s")275 print(f"[Startup] FLUX.2 Turbo LoRA path: {lora_path}")276 except Exception as ex:277 print(f"[Startup] FLUX.2 Turbo LoRA preload skipped/failed after {time.time() - t1:.2f}s: {ex}")278 279 280@torch.inference_mode()281def flux_pipe_call_that_returns_an_iterable_of_images(282 self,283 prompt: Union[str, List[str]] = None,284 prompt_2: Optional[Union[str, List[str]]] = None,285 negative_prompt: Optional[Union[str, List[str]]] = None,286 height: Optional[int] = None,287 width: Optional[int] = None,288 num_inference_steps: int = 28,289 timesteps: List[int] = None,290 guidance_scale: float = 3.5,291 num_images_per_prompt: Optional[int] = 1,292 generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,293 latents: Optional[torch.FloatTensor] = None,294 prompt_embeds: Optional[torch.FloatTensor] = None,295 pooled_prompt_embeds: Optional[torch.FloatTensor] = None,296 output_type: Optional[str] = "pil",297 return_dict: bool = True,298 joint_attention_kwargs: Optional[Dict[str, Any]] = None,299 max_sequence_length: int = 512,300 good_vae: Optional[Any] = None,301):302 # Set default height and width303 height = height or self.default_sample_size * self.vae_scale_factor304 width = width or self.default_sample_size * self.vae_scale_factor305 306 # Validate inputs307 self.check_inputs(308 prompt,309 prompt_2,310 height,311 width,312 prompt_embeds=prompt_embeds,313 pooled_prompt_embeds=pooled_prompt_embeds,314 max_sequence_length=max_sequence_length,315 )316 317 self._guidance_scale = guidance_scale318 self._joint_attention_kwargs = joint_attention_kwargs319 self._interrupt = False320 321 batch_size = 1 if isinstance(prompt, str) else len(prompt)322 device = self._execution_device323 324 lora_scale = joint_attention_kwargs.get("scale", None) if joint_attention_kwargs is not None else None325 326 # Prepare prompt inputs. Flux2Pipeline does not accept prompt_2 or negative_prompt327 # as separate kwargs, so merge them into a single prompt string for Flux2.328 if isinstance(self, Flux2Pipeline):329 combined_prompt = None330 if prompt is None:331 combined_prompt = prompt_2332 elif prompt_2 is None:333 combined_prompt = prompt334 else:335 combined_prompt = f"{prompt} {prompt_2}"336 prompt_to_encode = combined_prompt337 neg_prompt_to_encode = None338 else:339 prompt_to_encode = prompt340 neg_prompt_to_encode = None if negative_prompt is None else negative_prompt341 342 # Encode the positive prompt343 prompt_embeds_pos, pooled_prompt_embeds_pos, text_ids_pos = self.encode_prompt(344 prompt=prompt_to_encode,345 prompt_2=None,346 prompt_embeds=prompt_embeds,347 pooled_prompt_embeds=pooled_prompt_embeds,348 device=device,349 num_images_per_prompt=num_images_per_prompt,350 max_sequence_length=max_sequence_length,351 lora_scale=lora_scale,352 )353 354 # Encode the negative prompt if provided355 # Encode negative prompt only for pipelines that support it356 if (not isinstance(self, Flux2Pipeline)) and (negative_prompt is not None):357 prompt_embeds_neg, pooled_prompt_embeds_neg, text_ids_neg = self.encode_prompt(358 prompt=negative_prompt,359 prompt_2=None,360 prompt_embeds=None,361 pooled_prompt_embeds=None,362 device=device,363 num_images_per_prompt=num_images_per_prompt,364 max_sequence_length=max_sequence_length,365 lora_scale=lora_scale,366 )367 else:368 # Fallback to positive embeddings if no negative prompt or unsupported369 prompt_embeds_neg = prompt_embeds_pos370 pooled_prompt_embeds_neg = pooled_prompt_embeds_pos371 text_ids_neg = text_ids_pos372 373 # Prepare latents374 num_channels_latents = self.transformer.config.in_channels // 4375 latents, latent_image_ids = self.prepare_latents(376 batch_size * num_images_per_prompt,377 num_channels_latents,378 height,379 width,380 prompt_embeds_pos.dtype,381 device,382 generator,383 latents,384 )385 386 # Set up timesteps387 sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps)388 image_seq_len = latents.shape[1]389 mu = calculate_shift(390 image_seq_len,391 self.scheduler.config.base_image_seq_len,392 self.scheduler.config.max_image_seq_len,393 self.scheduler.config.base_shift,394 self.scheduler.config.max_shift,395 )396 timesteps, num_inference_steps = retrieve_timesteps(397 self.scheduler,398 num_inference_steps,399 device,400 timesteps,401 sigmas,402 mu=mu,403 )404 self._num_timesteps = len(timesteps)405 406 guidance = (407 torch.full([1], guidance_scale, device=device, dtype=torch.float32).expand(latents.shape[0])408 if self.transformer.config.guidance_embeds409 else None410 )411 412 # Denoising loop413 for i, t in enumerate(timesteps):414 if self._interrupt:415 continue416 417 timestep = t.expand(latents.shape[0]).to(latents.dtype)418 print(f"Step {i + 1}/{num_inference_steps} - Timestep: {timestep.item()}\n")419 420 # Compute noise prediction for positive prompt421 noise_pred_pos = self.transformer(422 hidden_states=latents,423 timestep=timestep / 1000,424 guidance=guidance,425 pooled_projections=pooled_prompt_embeds_pos,426 encoder_hidden_states=prompt_embeds_pos,427 txt_ids=text_ids_pos,428 img_ids=latent_image_ids,429 joint_attention_kwargs=self.joint_attention_kwargs,430 return_dict=False,431 )[0]432 433 # Compute noise prediction for negative prompt434 noise_pred_neg = self.transformer(435 hidden_states=latents,436 timestep=timestep / 1000,437 guidance=guidance,438 pooled_projections=pooled_prompt_embeds_neg,439 encoder_hidden_states=prompt_embeds_neg,440 txt_ids=text_ids_neg,441 img_ids=latent_image_ids,442 joint_attention_kwargs=self.joint_attention_kwargs,443 return_dict=False,444 )[0]445 446 # Combine noise predictions using guidance scale447 noise_pred = noise_pred_neg + guidance_scale * (noise_pred_pos - noise_pred_neg)448 449 # Generate intermediate image450 latents_for_image = self._unpack_latents(latents, height, width, self.vae_scale_factor)451 latents_for_image = (latents_for_image / self.vae.config.scaling_factor) + self.vae.config.shift_factor452 image = self.vae.decode(latents_for_image, return_dict=False)[0]453 yield self.image_processor.postprocess(image, output_type=output_type)[0]454 455 # Update latents with combined noise prediction456 latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]457 torch.cuda.empty_cache()458 459 # Final image generation460 latents = self._unpack_latents(latents, height, width, self.vae_scale_factor)461 latents = (latents / good_vae.config.scaling_factor) + good_vae.config.shift_factor462 image = good_vae.decode(latents, return_dict=False)[0]463 self.maybe_free_model_hooks()464 torch.cuda.empty_cache()465 yield self.image_processor.postprocess(image, output_type=output_type)[0]466 467 468@torch.inference_mode()469def flux_pipe_call_single_final_image(470 self,471 prompt: Union[str, List[str]] = None,472 prompt_2: Optional[Union[str, List[str]]] = None,473 negative_prompt: Optional[Union[str, List[str]]] = None,474 height: Optional[int] = None,475 width: Optional[int] = None,476 num_inference_steps: int = 28,477 guidance_scale: float = 3.5,478 generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,479 joint_attention_kwargs: Optional[Dict[str, Any]] = None,480 output_type: Optional[str] = "pil",481):482 is_f2 = isinstance(self, Flux2Pipeline)483 final_steps = num_inference_steps if num_inference_steps is not None else (7 if is_f2 else 28)484 485 kwargs = {486 "prompt": prompt,487 "height": height,488 "width": width,489 "num_inference_steps": final_steps,490 "guidance_scale": guidance_scale,491 "generator": generator,492 "output_type": output_type,493 }494 495 if is_f2:496 # Flux2Pipeline does not accept prompt_2; merge into prompt.497 if prompt is None:498 kwargs["prompt"] = prompt_2499 elif prompt_2 is not None:500 kwargs["prompt"] = f"{prompt} {prompt_2}"501 else:502 if prompt_2 is not None:503 kwargs["prompt_2"] = prompt_2504 505 # Flux2Pipeline does not support negative_prompt506 if (negative_prompt is not None) and (not is_f2):507 kwargs["negative_prompt"] = negative_prompt508 509 if joint_attention_kwargs is not None:510 if is_f2:511 kwargs["attention_kwargs"] = joint_attention_kwargs512 else:513 kwargs["joint_attention_kwargs"] = joint_attention_kwargs514 515 return self(**kwargs).images[0]516#--------------------------------------------------Model Initialization-----------------------------------------------------------------------------------------#517 518# Pre-download FLUX.2 once at startup so later switching is cache-hit.519preload_flux2_model_to_cache()520 521dtype = torch.bfloat16522device = "cpu" # ZeroGPU-safe at startup; move to CUDA only inside @spaces.GPU523base_model = "black-forest-labs/FLUX.1-dev"524active_base_model = base_model525flux2_cpu_offload_enabled = False526 527taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype)528good_vae = AutoencoderKL.from_pretrained(base_model, subfolder="vae", torch_dtype=dtype)529pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=dtype, vae=taef1)530pipe_i2i = None531# pipe_i2i = AutoPipelineForImage2Image.from_pretrained(base_model,532# vae=good_vae,533# transformer=pipe.transformer,534# text_encoder=pipe.text_encoder,535# tokenizer=pipe.tokenizer,536# text_encoder_2=pipe.text_encoder_2,537# tokenizer_2=pipe.tokenizer_2,538# torch_dtype=dtype539# )540 541pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)542pipe.flux_pipe_call_single_final_image = flux_pipe_call_single_final_image.__get__(pipe)543 544def _release_gpu_obj(obj):545 if obj is None:546 return None547 try:548 if hasattr(obj, "to"):549 obj.to("cpu")550 except Exception:551 pass552 del obj553 gc.collect()554 # IMPORTANT: no torch.cuda.* calls here (can run outside ZeroGPU allowance)555 return None556 557 558def ensure_pipelines_for_model(model_name: str, use_safetensors: bool = False, image_input=None):559 global pipe, pipe_i2i, good_vae, active_base_model, flux2_cpu_offload_enabled560 need_i2i = image_input is not None561 already_correct_model = (model_name == active_base_model)562 563 # Fast path: correct model and requested variant already exists.564 if already_correct_model:565 if need_i2i and pipe_i2i is not None:566 pipe = _release_gpu_obj(pipe) # enforce exclusivity567 status_msg = f"Using base model: {active_base_model} (i2i ready)"568 print(status_msg)569 gr.Info(status_msg)570 return571 if (not need_i2i) and pipe is not None:572 pipe_i2i = _release_gpu_obj(pipe_i2i) # enforce exclusivity573 status_msg = f"Using base model: {active_base_model} (t2i ready)"574 print(status_msg)575 gr.Info(status_msg)576 return577 578 status_msg = f"Rebuilding {'i2i' if need_i2i else 't2i'} pipeline for {model_name}"579 print(status_msg)580 gr.Info(status_msg)581 else:582 status_msg = f"Switching base model from {active_base_model} to {model_name}"583 print(status_msg)584 gr.Info(status_msg)585 # release old model resources before rebuilding586 pipe = _release_gpu_obj(pipe)587 pipe_i2i = _release_gpu_obj(pipe_i2i)588 good_vae = _release_gpu_obj(good_vae)589 590 # Build t2i source pipe if needed (also needed as source when creating i2i).591 if pipe is None:592 if is_flux2_model(model_name):593 text_encoder = Mistral3ForConditionalGeneration.from_pretrained(594 model_name,595 subfolder="text_encoder",596 torch_dtype=torch.bfloat16,597 device_map="cpu",598 local_files_only=True,599 )600 601 dit = AutoModel.from_pretrained(602 model_name,603 subfolder="transformer",604 torch_dtype=torch.bfloat16,605 device_map="cpu",606 local_files_only=True,607 )608 609 pipe = Flux2Pipeline.from_pretrained(610 model_name,611 torch_dtype=dtype,612 text_encoder=text_encoder,613 transformer=dit,614 use_safetensors=use_safetensors,615 local_files_only=True,616 )617 flux2_cpu_offload_enabled = should_enable_flux2_cpu_offload(60.0)618 if flux2_cpu_offload_enabled:619 pipe.enable_model_cpu_offload()620 print(f"FLUX.2 CPU offload enabled (VRAM: {get_total_vram_gb():.2f} GB)")621 else:622 print(f"FLUX.2 CPU offload disabled (VRAM: {get_total_vram_gb():.2f} GB)")623 good_vae = pipe.vae624 elif "FLUX" in model_name.upper():625 flux2_cpu_offload_enabled = False626 good_vae = AutoencoderKL.from_pretrained(627 model_name, subfolder="vae", torch_dtype=dtype, use_safetensors=use_safetensors628 )629 pipe = FluxPipeline.from_pretrained(630 model_name, torch_dtype=dtype, vae=taef1, use_safetensors=use_safetensors631 )632 else:633 flux2_cpu_offload_enabled = False634 good_vae = AutoencoderKL.from_pretrained(635 model_name, subfolder="vae", torch_dtype=dtype, use_safetensors=use_safetensors636 )637 pipe = DiffusionPipeline.from_pretrained(638 model_name, torch_dtype=dtype, vae=taef1, use_safetensors=use_safetensors639 )640 641 pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)642 pipe.flux_pipe_call_single_final_image = flux_pipe_call_single_final_image.__get__(pipe)643 644 if need_i2i:645 pipe_i2i = _release_gpu_obj(pipe_i2i)646 647 if is_flux2_model(model_name):648 # FLUX.2 i2i: reuse Flux2Pipeline (no AutoPipelineForImage2Image)649 pipe_i2i = pipe650 pipe = None651 else:652 pipe_i2i_kwargs = {653 "vae": good_vae,654 "transformer": pipe.transformer,655 "torch_dtype": dtype,656 "use_safetensors": use_safetensors,657 "local_files_only": True,658 }659 if hasattr(pipe, "text_encoder") and pipe.text_encoder is not None:660 pipe_i2i_kwargs["text_encoder"] = pipe.text_encoder661 if hasattr(pipe, "tokenizer") and pipe.tokenizer is not None:662 pipe_i2i_kwargs["tokenizer"] = pipe.tokenizer663 if hasattr(pipe, "text_encoder_2") and pipe.text_encoder_2 is not None:664 pipe_i2i_kwargs["text_encoder_2"] = pipe.text_encoder_2665 if hasattr(pipe, "tokenizer_2") and pipe.tokenizer_2 is not None:666 pipe_i2i_kwargs["tokenizer_2"] = pipe.tokenizer_2667 668 pipe_i2i = AutoPipelineForImage2Image.from_pretrained(model_name, **pipe_i2i_kwargs)669 pipe = _release_gpu_obj(pipe)670 else:671 pipe_i2i = _release_gpu_obj(pipe_i2i)672 673 active_base_model = model_name674 675def _move_for_inference(main_pipe, vae_obj):676 if isinstance(main_pipe, Flux2Pipeline) and flux2_cpu_offload_enabled:677 return678 if torch.cuda.is_available():679 try:680 main_pipe.to("cuda")681 vae_obj.to("cuda")682 except RuntimeError as ex:683 if "CUDACachingAllocator" in str(ex) or "NVML_SUCCESS" in str(ex):684 torch.cuda.empty_cache()685 try:686 torch.cuda.ipc_collect()687 except Exception:688 pass689 main_pipe.to("cuda")690 vae_obj.to("cuda")691 else:692 raise693 694class calculateDuration:695 def __init__(self, activity_name=""):696 self.activity_name = activity_name697 698 def __enter__(self):699 self.start_time = time.time()700 return self701 702 def __exit__(self, exc_type, exc_value, traceback):703 self.end_time = time.time()704 self.elapsed_time = self.end_time - self.start_time705 if self.activity_name:706 print(f"Elapsed time for {self.activity_name}: {self.elapsed_time:.6f} seconds")707 else:708 print(f"Elapsed time: {self.elapsed_time:.6f} seconds")709 710def update_selection(evt: gr.SelectData, width, height, aspect_ratio, current_steps):711 selected_lora = loras[evt.index]712 new_placeholder = f"Type a prompt for {selected_lora['title']}"713 new_aspect_ratio = aspect_ratio714 lora_repo = selected_lora["repo"]715 target_model = LORA_TO_MODEL.get(lora_repo, base_model)716 lora_runtime_parameters = get_lora_runtime_parameters(lora_repo, selected_lora)717 lora_default_steps = lora_runtime_parameters.get("num_inference_steps")718 if isinstance(lora_default_steps, float):719 lora_default_steps = int(lora_default_steps)720 721 if is_flux2_model(target_model):722 flux2_steps = lora_default_steps if isinstance(lora_default_steps, int) else 7723 steps_update = gr.update(724 minimum=flux2_steps,725 maximum=flux2_steps,726 step=1,727 value=flux2_steps,728 interactive=False,729 label="Steps (from LoRA defaults for FLUX.2)",730 )731 else:732 safe_steps = 7733 try:734 if current_steps is not None:735 safe_steps = int(current_steps)736 except Exception:737 safe_steps = 7738 steps_update = gr.update(739 minimum=1,740 maximum=50,741 step=1,742 value=max(1, min(50, safe_steps)),743 interactive=True,744 label="Steps",745 )746 updated_text = f"### Selected: [{lora_repo}](https://huggingface.co/{lora_repo}) ✅"747 # aspect will now use ratios if implemented, like 16:9, 4:3, 1:1748 if "aspect" in selected_lora:749 try:750 new_aspect_ratio = selected_lora["aspect"]751 width, height = update_dimensions_on_ratio(new_aspect_ratio, height)752 except Exception as e:753 print(f"\nError in update selection aspect ratios:{e}\nSkipping")754 new_aspect_ratio = aspect_ratio755 width = width756 height = height757 return (758 gr.update(placeholder=new_placeholder),759 updated_text,760 evt.index,761 width,762 height,763 new_aspect_ratio,764 upd_prompt_notes_by_index(evt.index),765 steps_update,766 )767 768@spaces.GPU(duration=120,progress=gr.Progress(track_tqdm=True))769def generate_image(prompt_mash, negative_prompt, steps, seed, cfg_scale, width, height, lora_scale, progress):770 _move_for_inference(pipe, good_vae)771 generator = torch.Generator(device="cuda").manual_seed(seed)772 flash_attention_enabled = torch.backends.cuda.flash_sdp_enabled()773 if flash_attention_enabled:774 pipe.attn_implementation="flash_attention_2"775 776 if IS_SHARED_SPACE:777 pipe.vae.enable_tiling() # For larger resolutions if needed778 else:779 # Compile UNet780 #pipe.transformer = torch.compile(pipe.transformer, mode="reduce-overhead")781 #pipe.enable_model_cpu_offload() #for smaller GPUs782 pipe.vae.enable_slicing()783 784 # Disable unnecessary features785 pipe.safety_checker = None786 print(f"\nGenerating image with prompt: {prompt_mash}\n")787 approx_tokens= approximate_token_count(prompt_mash)788 if approx_tokens > 76:789 print(f"\nSplitting prompt due to length: {approx_tokens}\n")790 prompt, prompt2 = split_prompt_precisely(prompt_mash)791 else:792 prompt = prompt_mash793 prompt2 = None794 with calculateDuration("Generating image"):795 # Generate image796 if isinstance(pipe, Flux2Pipeline):797 final_image = pipe.flux_pipe_call_single_final_image(798 prompt=prompt,799 prompt_2=prompt2, 800 num_inference_steps=steps,801 guidance_scale=cfg_scale,802 width=width,803 height=height,804 generator=generator,805 joint_attention_kwargs={"scale": lora_scale},806 output_type="pil",807 )808 yield final_image809 else:810 for img in pipe.flux_pipe_call_that_returns_an_iterable_of_images(811 prompt=prompt,812 prompt_2=prompt2,813 num_inference_steps=steps,814 guidance_scale=cfg_scale,815 width=width,816 height=height,817 generator=generator,818 joint_attention_kwargs={"scale": lora_scale},819 output_type="pil",820 good_vae=good_vae,821 ):822 yield img823 824@spaces.GPU(duration=120,progress=gr.Progress(track_tqdm=True))825def generate_image_to_image(prompt_mash, negative_prompt, image_input_path, image_strength, steps, cfg_scale, width, height, lora_scale, seed, progress):826 _move_for_inference(pipe_i2i, good_vae)827 generator = torch.Generator(device="cuda").manual_seed(seed) 828 flash_attention_enabled = torch.backends.cuda.flash_sdp_enabled()829 if flash_attention_enabled:830 pipe_i2i.attn_implementation="flash_attention_2"831 832 if IS_SHARED_SPACE:833 pipe_i2i.vae.enable_tiling() # For larger resolutions if needed834 else:835 # Compile UNet836 # pipe.transformer = torch.compile(pipe.transformer, mode="reduce-overhead") # uses the other pipe's transformer837 # pipe_i2i.enable_model_cpu_offload() #for smaller GPUs838 pipe_i2i.vae.enable_slicing()839 840 841 # Disable unnecessary features842 pipe_i2i.safety_checker = None843 image_input = open_image(image_input_path)844 print(f"\nGenerating image with prompt: {prompt_mash} and {image_input_path}\n")845 approx_tokens= approximate_token_count(prompt_mash)846 if approx_tokens > 76:847 print(f"\nSplitting prompt due to length: {approx_tokens}\n")848 prompt, prompt2 = split_prompt_precisely(prompt_mash)849 else:850 prompt = prompt_mash851 prompt2 = None852 with calculateDuration("Generating image"):853 # Generate image854 if isinstance(pipe_i2i, Flux2Pipeline):855 # Flux2Pipeline does not accept prompt_2; merge prompt and prompt_2856 if prompt is None:857 merged_prompt = prompt2858 elif prompt2 is None:859 merged_prompt = prompt860 else:861 merged_prompt = f"{prompt} {prompt2}"862 863 final_image = pipe_i2i(864 prompt=merged_prompt,865 image=image_input,866 num_inference_steps=steps,867 guidance_scale=cfg_scale,868 width=width,869 height=height,870 generator=generator,871 attention_kwargs={"scale": lora_scale},872 output_type="pil",873 ).images[0]874 else:875 final_image = pipe_i2i(876 prompt=prompt,877 prompt_2=prompt2,878 image=image_input,879 strength=image_strength,880 num_inference_steps=steps,881 guidance_scale=cfg_scale,882 width=width,883 height=height,884 generator=generator,885 joint_attention_kwargs={"scale": lora_scale},886 output_type="pil",887 ).images[0]888 return final_image889 890 891def 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]]:892 if selected_index is None:893 raise gr.Error("You must select a LoRA before proceeding.🧨")894 print(f"input Image: {image_input}\n")895 # handle selecting a conditioned image from the gallery896 global current_prerendered_image897 conditioned_image=None898 formatted_map_option = map_option.lower().replace(' ', '_')899 negative_prompt = ""900 901 if use_conditioned_image:902 print(f"Conditioned path: {current_prerendered_image.value}.. converting to RGB\n")903 # ensure the conditioned image is an image and not a string, cannot use RGBA904 if isinstance(current_prerendered_image.value, str):905 conditioned_image = open_image(current_prerendered_image.value).convert("RGB")906 image_input = resize_and_crop_image(conditioned_image, width, height)907 print(f"Conditioned Image: {image_input.size}.. converted to RGB and resized\n")908 if map_option != "Prompt":909 prompt = PROMPTS[map_option]910 negative_prompt = NEGATIVE_PROMPTS.get(map_option, "")911 912 selected_lora = loras[selected_index]913 lora_path = selected_lora["repo"]914 target_model = LORA_TO_MODEL.get(lora_path, base_model)915 lora_runtime_parameters = get_lora_runtime_parameters(lora_path, selected_lora)916 917 # detect safetensor usage918 use_safetensors_flag = False919 # check explicit "weights" field on the lora entry920 weight_field = selected_lora.get("weights") or ""921 if weight_field and "safetensor" in weight_field.lower():922 use_safetensors_flag = True923 else:924 # fallback: check LORA_DETAILS for weight_name entries925 details = LORA_DETAILS.get(lora_path)926 if details:927 for d in details:928 wn = d.get("weight_name") or d.get("weights") or ""929 if wn and "safetensor" in wn.lower():930 use_safetensors_flag = True931 break932 933 ensure_pipelines_for_model(target_model, use_safetensors=use_safetensors_flag, image_input=image_input)934 935 runtime_steps = lora_runtime_parameters.get("num_inference_steps")936 runtime_guidance = lora_runtime_parameters.get("guidance_scale")937 938 if isinstance(runtime_steps, (int, float)):939 steps = int(runtime_steps)940 elif is_flux2_model(active_base_model):941 steps = 8942 943 if isinstance(runtime_guidance, (int, float)):944 cfg_scale = float(runtime_guidance)945 trigger_word = selected_lora["trigger_word"]946 if(trigger_word):947 if "trigger_position" in selected_lora:948 if selected_lora["trigger_position"] == "prepend":949 prompt_mash = f"{trigger_word} {prompt}"950 else:951 prompt_mash = f"{prompt} {trigger_word}"952 else:953 prompt_mash = f"{trigger_word} {prompt}"954 else:955 prompt_mash = prompt956 957 with calculateDuration("Unloading LoRA"):958 if pipe is not None:959 pipe.unload_lora_weights()960 if pipe_i2i is not None:961 pipe_i2i.unload_lora_weights()962 963 #LoRA weights flow964 with calculateDuration(f"Loading LoRA weights for {selected_lora['title']}"):965 pipe_to_use = pipe_i2i if image_input is not None else pipe966 weight_name = selected_lora.get("weights", None)967 968 pipe_to_use.load_lora_weights(969 lora_path, 970 weight_name=weight_name, 971 low_cpu_mem_usage=True972 )973 974 with calculateDuration("Randomizing seed"):975 if randomize_seed:976 seed = random.randint(0, MAX_SEED)977 978 if(image_input is not None):979 print(f"\nGenerating image to image with seed: {seed}\n")980 generated_image = generate_image_to_image(prompt_mash, negative_prompt, image_input, image_strength, steps, cfg_scale, width, height, lora_scale, seed, progress)981 982 if enlarge:983 upscale_factor = max(1.0, min((TARGET_SIZE[0] / width), (TARGET_SIZE[1] / height)))984 progress(0.75, desc=f"upscaling to {upscale_factor}")985 upscaled_image = upscale_image(generated_image, upscale_factor)986 progress(1.0, desc=f"upscaled to {upscale_factor}")987 else:988 upscaled_image = generated_image989 # Save the upscaled image to a temporary file990 with NamedTemporaryFile(delete=False, suffix=".png", prefix=f"{formatted_map_option}_") as tmp_upscaled:991 upscaled_image.save(tmp_upscaled.name, format="PNG")992 temp_files.append(tmp_upscaled.name)993 print(f"Upscaled image saved to {tmp_upscaled.name}")994 final_image = tmp_upscaled.name995 yield final_image, seed, gr.update(visible=False)996 else:997 image_generator = generate_image(prompt_mash, negative_prompt, steps, seed, cfg_scale, width, height, lora_scale, progress)998 999 final_image = None1000 step_counter = 01001 # create a single temporary file for intermediate previews1002 with NamedTemporaryFile(delete=False, suffix=".png", prefix=f"{formatted_map_option}_step_") as tmp_step:1003 tmp_preview_path = tmp_step.name1004 1005 for image in image_generator:1006 step_counter += 11007 generated_image = image # PIL.Image.Image1008 # overwrite the same preview file to update UI without creating many files1009 try:1010 generated_image.save(tmp_preview_path, format="PNG")1011 except Exception:1012 # fallback: save to a new temp file if overwrite fails1013 with NamedTemporaryFile(delete=False, suffix=".png", prefix=f"{formatted_map_option}_step_") as tf:1014 generated_image.save(tf.name, format="PNG")1015 tmp_preview_path = tf.name1016 1017 progress_bar = f'<div class="progress-container"><div class="progress-bar" style="--current: {step_counter}; --total: {steps};"></div></div>'1018 # yield preview file path, seed, progress1019 yield tmp_preview_path, seed, gr.update(value=progress_bar, visible=True)1020 1021 if enlarge:1022 upscale_factor = max(1.0, min((TARGET_SIZE[0] / width), (TARGET_SIZE[1] / height)))1023 progress(0.98, desc=f"upscaling to {upscale_factor}")1024 upscaled_image = upscale_image(generated_image, upscale_factor)1025 else:1026 upscaled_image = generated_image1027 # Save the upscaled image to a temporary file1028 with NamedTemporaryFile(delete=False, suffix=".png", prefix=f"{formatted_map_option}_") as tmp_upscaled:1029 upscaled_image.save(tmp_upscaled.name, format="PNG")1030 temp_files.append(tmp_upscaled.name)1031 print(f"Upscaled image saved to {tmp_upscaled.name}")1032 final_image = tmp_upscaled.name1033 yield final_image, seed, gr.update(value=progress_bar, visible=False)1034 1035def get_huggingface_safetensors(link):1036 split_link = link.split("/")1037 if len(split_link) == 2:1038 model_card = ModelCard.load(link)1039 base_model = model_card.data.get("base_model")1040 print(base_model)1041 1042 # Allows Both1043 if base_model not in MODELS:1044 raise Exception("Flux LoRA Not Found!")1045 1046 image_path = model_card.data.get("widget", [{}])[0].get("output", {}).get("url", None)1047 trigger_word = model_card.data.get("instance_prompt", "")1048 image_url = f"https://huggingface.co/{link}/resolve/main/{image_path}" if image_path else None1049 fs = HfFileSystem()1050 1051 safetensors_name = None1052 try:1053 list_of_files = fs.ls(link, detail=False)1054 for file in list_of_files:1055 if file.endswith(".safetensors"):1056 safetensors_name = file.split("/")[-1]1057 if (not image_url) and file.lower().endswith((".jpg", ".jpeg", ".png", ".webp")):1058 image_elements = file.split("/")1059 image_url = f"https://huggingface.co/{link}/resolve/main/{image_elements[-1]}"1060 except Exception as e:1061 print(e)1062 gr.Warning("You didn't include a link neither a valid Hugging Face repository with a *.safetensors LoRA")1063 raise Exception("You didn't include a link neither a valid Hugging Face repository with a *.safetensors LoRA")1064 1065 return split_link[1], link, safetensors_name, trigger_word, image_url, base_model1066 1067def check_custom_model(link):1068 if(link.startswith("https://")):1069 if(link.startswith("https://huggingface.co") or link.startswith("https://www.huggingface.co")):1070 link_split = link.split("huggingface.co/")1071 return get_huggingface_safetensors(link_split[1])1072 else: 1073 return get_huggingface_safetensors(link)1074 1075def add_custom_lora(custom_lora):1076 global loras1077 if(custom_lora):1078 try:1079 title, repo, path, trigger_word, image, detected_base_model = check_custom_model(custom_lora)1080 print(f"Loaded custom LoRA: {repo}")1081 card = f'''1082 <div class="custom_lora_card">1083 <span>Loaded custom LoRA:</span>1084 <div class="card_internal">1085 <img src="{image}" />1086 <div>1087 <h3>{title}</h3>1088 <small>{"Using: <code><b>"+trigger_word+"</code></b> as the trigger word" if trigger_word else "No trigger word found. If there's a trigger word, include it in your prompt"}<br></small>1089 </div>1090 </div>1091 </div>1092 '''1093 existing_item_index = next((index for (index, item) in enumerate(loras) if item['repo'] == repo), None)1094 if(not existing_item_index):1095 new_item = {1096 "image": image,1097 "title": title,1098 "repo": repo,1099 "weights": path,1100 "trigger_word": trigger_word,1101 "base_model": detected_base_model1102 }1103 print(new_item)1104 existing_item_index = len(loras)1105 loras.append(new_item)1106 1107 return gr.update(visible=True, value=card), gr.update(visible=True), gr.Gallery(selected_index=None), f"Custom: {path}", existing_item_index, trigger_word1108 except Exception as e:1109 gr.Warning(f"Invalid LoRA: either you entered an invalid link, or a non-FLUX LoRA")1110 return gr.update(visible=True, value=f"Invalid LoRA: either you entered an invalid link, a non-FLUX LoRA"), gr.update(visible=False), gr.update(), "", None, ""1111 else:1112 return gr.update(visible=False), gr.update(visible=False), gr.update(), "", None, ""1113 1114def remove_custom_lora():1115 return gr.update(visible=False), gr.update(visible=False), gr.update(), "", None, ""1116 1117def on_prerendered_gallery_selection(event_data: gr.SelectData):1118 global current_prerendered_image1119 selected_index = event_data.index1120 selected_image = pre_rendered_maps_paths[selected_index]1121 print(f"Gallery Image Selected: {selected_image}\n")1122 current_prerendered_image.value = selected_image1123 return current_prerendered_image1124 1125def update_prompt_visibility(map_option):1126 is_visible = (map_option == "Prompt")1127 return (1128 gr.update(visible=is_visible),1129 gr.update(visible=is_visible),1130 gr.update(visible=is_visible)1131 )1132 1133def copy_map_prompts_to_manual(map_option: str):1134 """Return the manual prompt and negative prompt values for the given map option."""1135 # Safely fetch values1136 prompt_val = PROMPTS.get(map_option) if map_option is not None else None1137 negative_val = NEGATIVE_PROMPTS.get(map_option) if map_option is not None else None1138 1139 # If a PROMPTS/NEGATIVE_PROMPTS entry is None, do not overwrite that textbox1140 # Return gr.update() (no-op) for components that should remain unchanged.1141 out_prompt = gr.update(value=prompt_val) if prompt_val is not None else gr.update()1142 out_negative = gr.update(value=negative_val) if negative_val is not None else gr.update()1143 return out_prompt, out_negative1144def composite_with_control_sync(input_image, sketch_image, slider_value):1145 if input_image is None or sketch_image is None:1146 return input_image1147 1148 # Load the images using open_image() if they are provided as file paths.1149 in_img = open_image(input_image) if isinstance(input_image, str) else input_image1150 if in_img is None:1151 return input_image1152 1153 sk_img_path, _ = get_image_from_dict(sketch_image)1154 if not sk_img_path:1155 return input_image1156 1157 sk_img = open_image(sk_img_path)1158 if sk_img is None:1159 return input_image1160 1161 # Resize sketch image if dimensions don't match input image.1162 if in_img.size != sk_img.size:1163 sk_img = sk_img.resize(in_img.size, Image.LANCZOS)1164 1165 # Now composite using the original alpha_composite_with_control function.1166 result_img = alpha_composite_with_control(in_img, sk_img, slider_value)1167 return result_img1168 1169def replace_input_with_sketch_image(sketch_image):1170 print(f"Sketch Image: {sketch_image}\n")1171 sketch, is_dict = get_image_from_dict(sketch_image)1172 if sketch is None:1173 gr.Warning("Sketch image not found.")1174 return None1175 try:1176 sketch_img = open_image(sketch) if isinstance(sketch, str) else sketch1177 with NamedTemporaryFile(delete=False, suffix=".png", prefix="sketch_input_") as tmp_sketch:1178 sketch_img.convert("RGBA").save(tmp_sketch.name, format="PNG")1179 temp_files.append(tmp_sketch.name)1180 return tmp_sketch.name1181 except Exception as e:1182 print(f"Error preparing sketch image for input: {e}")1183 return sketch1184 1185def on_input_image_change(image_path):1186 if image_path is None:1187 gr.Warning("Please upload an Input Image to get started.")1188 return None, gr.update()1189 img, img_path = convert_to_rgba_png(image_path)1190 with Image.open(img_path) as pil_img:1191 width, height = pil_img.size1192 return [img_path, gr.update(width=width, height=height)]1193 1194def update_sketch_dimensions(input_image, sketch_image):1195 if input_image is None or sketch_image is None:1196 return [sketch_image, gr.update()]1197 1198 # Load the images using open_image() if they are provided as file paths.1199 in_img = open_image(input_image) if isinstance(input_image, str) else input_image1200 if in_img is None: