import spaces import gradio as gr from huggingface_hub import ModelCard from config import Config from .models import * from .handlers import gen_img # Common def update_model_options(model): for m in Config.IMAGES_MODELS: if m['repo_id'] == model: if m['loader'] == 'flux': return ( gr.update( # negative_prompt visible=False ), gr.update( # lora_gallery value=[(lora['image'], lora['title']) for lora in Config.IMAGES_LORAS_FLUX] ), gr.update( # embeddings_accordion visible=False ), gr.update( # scribble_tab visible=False ), gr.update( # scheduler value='fm_euler' ), gr.update( # image_clip_skip visible=False ), gr.update( # image_guidance_scale value=3.5 ) ) elif m['loader'] == 'sdxl': return ( gr.update( # negative_prompt visible=True ), gr.update( # lora_gallery value=[(lora['image'], lora['title']) for lora in Config.IMAGES_LORAS_SDXL] ), gr.update( # embeddings_accordion visible=True ), gr.update( # scribble_tab visible=True ), gr.update( # scheduler value='dpmpp_2m_sde_k' ), gr.update( # image_clip_skip visible=True ), gr.update( # image_guidance_scale value=7.0 ) ) def update_fast_generation(model, fast_generation): for m in Config.IMAGES_MODELS: if m['repo_id'] == model: if m['loader'] == 'flux': if fast_generation: return ( gr.update( # image_num_inference_steps value=8 ), gr.update( # image_guidance_scale value=3.5 ) ) else: return ( gr.update( # image_num_inference_steps value=20 ), gr.update( # image_guidance_scale value=3.5 ) ) elif m['loader'] == 'sdxl': if fast_generation: return ( gr.update( # image_num_inference_steps value=8 ), gr.update( # image_guidance_scale value=1.0 ) ) else: return ( gr.update( # image_num_inference_steps value=20 ), gr.update( # image_guidance_scale value=7.0 ) ) # Loras def selected_lora_from_gallery(evt: gr.SelectData): return ( gr.update( value=evt.index ) ) def update_selected_lora(custom_lora): link = custom_lora.split("/") if len(link) == 2: model_card = ModelCard.load(custom_lora) trigger_word = model_card.data.get("instance_prompt", "") image_url = f"""https://huggingface.co/{custom_lora}/resolve/main/{model_card.data.get("widget", [{}])[0].get("output", {}).get("url", None)}""" custom_lora_info_css = """ """ custom_lora_info_html = f"""
Custom LoRA: {custom_lora}
LoRA preview

{link[1].replace("-", " ").replace("_", " ")}

{"Using: "+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"}
""" custom_lora_info_html = f"{custom_lora_info_css}{custom_lora_info_html}" return ( gr.update( # selected_lora value=custom_lora, ), gr.update( # custom_lora_info value=custom_lora_info_html, visible=True ) ) else: return ( gr.update( # selected_lora value=custom_lora, ), gr.update( # custom_lora_info value=custom_lora_info_html if len(link) == 0 else "", visible=False ) ) def update_lora_sliders(enabled_loras): sliders = [] remove_buttons = [] for lora in enabled_loras: sliders.append( gr.update( label=lora.get("repo_id", ""), info=f"Trigger Word: {lora.get('trigger_word', '')}", visible=True, interactive=True ) ) remove_buttons.append( gr.update( visible=True, interactive=True ) ) if len(sliders) < 6: for i in range(len(sliders), 6): sliders.append( gr.update( visible=False ) ) remove_buttons.append( gr.update( visible=False ) ) return *sliders, *remove_buttons def remove_from_enabled_loras(enabled_loras, index): enabled_loras.pop(index) return ( gr.update( value=enabled_loras ) ) def add_to_enabled_loras(model, selected_lora, enabled_loras): for m in Config.IMAGES_MODELS: if m['repo_id'] == model: lora_data = [] if m['loader'] == 'flux': lora_data = Config.IMAGES_LORAS_FLUX elif m['loader'] == 'sdxl': lora_data = Config.IMAGES_LORAS_SDXL try: selected_lora = int(selected_lora) if 0 <= selected_lora: # is the index of the lora in the gallery lora_info = lora_data[selected_lora] enabled_loras.append({ "repo_id": lora_info["repo"], "trigger_word": lora_info["trigger_word"] }) except ValueError: link = selected_lora.split("/") if len(link) == 2: model_card = ModelCard.load(selected_lora) trigger_word = model_card.data.get("instance_prompt", "") enabled_loras.append({ "repo_id": selected_lora, "trigger_word": trigger_word }) return ( gr.update( # selected_lora value="" ), gr.update( # custom_lora_info value="", visible=False ), gr.update( # enabled_loras value=enabled_loras ) ) # Custom Embedding def update_custom_embedding(custom_embedding): link = custom_embedding.split("/") if len(link) == 2: model_card = ModelCard.load(custom_embedding) trigger_word = model_card.data.get("instance_prompt", "") image_url = f"""https://huggingface.co/{custom_embedding}/resolve/main/{model_card.data.get("widget", [{}])[0].get("output", {}).get("url", None)}""" custom_embedding_info_css = """ """ custom_embedding_info_html = f"""
Custom Embedding: {custom_embedding}
Embedding preview

{link[1].replace("-", " ").replace("_", " ")}

{"Using: "+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"}
""" custom_embedding_info_html = f"{custom_embedding_info_css}{custom_embedding_info_html}" return gr.update(value=custom_embedding_info_html, visible=True) else: return gr.update(value="", visible=False) def add_to_embeddings(custom_embedding, enabled_embeddings): link = custom_embedding.split("/") if len(link) == 2: if ModelCard.load(custom_embedding): enabled_embeddings.append(custom_embedding) return ( gr.update( # custom_embedding value="" ), gr.update( # custom_embedding_info value="", visible=False ), gr.update( # enabled_embeddings value=enabled_embeddings ) ) def update_enabled_embeddings_list(enabled_embeddings): return gr.update( # enabled_embeddings_list value=enabled_embeddings, choices=enabled_embeddings ) def update_enabled_embeddings(enabled_embeddings_list): return gr.update( # enabled_embeddings value=enabled_embeddings_list ) # Generate Image @spaces.GPU(duration=75) def generate_image( model, prompt, negative_prompt, fast_generation, enabled_loras, enabled_embeddings, lora_slider_0, lora_slider_1, lora_slider_2, lora_slider_3, lora_slider_4, lora_slider_5, # type: ignore img2img_image, inpaint_image, canny_image, pose_image, depth_image, scribble_image, # type: ignore img2img_strength, inpaint_strength, canny_strength, pose_strength, depth_strength, scribble_strength, # type: ignore resize_mode, scheduler, image_height, image_width, image_num_images_per_prompt, # type: ignore image_num_inference_steps, image_clip_skip, image_guidance_scale, image_seed, # type: ignore refiner, vae, progress=gr.Progress(track_tqdm=True) ): try: progress(0, "Configuring arguments...") base_args = { "model": model, "prompt": prompt, # "negative_prompt": negative_prompt, "fast_generation": fast_generation, "loras": None, # "embeddings": None, "resize_mode": resize_mode, "scheduler": scheduler, "height": int(image_height), "width": int(image_width), "num_images_per_prompt": float(image_num_images_per_prompt), "num_inference_steps": float(image_num_inference_steps), # "clip_skip": None, "guidance_scale": image_guidance_scale, "seed": int(image_seed), "refiner": refiner, "vae": vae, "controlnet_config": None, } base_args = BaseReq(**base_args) if len(enabled_loras) > 0: base_args.loras = [] for enabled_lora, slider in zip(enabled_loras, [lora_slider_0, lora_slider_1, lora_slider_2, lora_slider_3, lora_slider_4, lora_slider_5]): if enabled_lora['repo_id']: base_args.loras.append({ "repo_id": enabled_lora['repo_id'], "weight": slider }) # Load SDXL related args if model in Config.IMAGES_MODELS: if model['loader'] == 'sdxl': base_args.negative_prompt = negative_prompt base_args.clip_skip = image_clip_skip if len(enabled_embeddings) > 0: base_args.embeddings = enabled_embeddings image = None mask_image = None strength = None if img2img_image: image = img2img_image strength = float(img2img_strength) base_args = BaseImg2ImgReq( **base_args.__dict__, image=image, strength=strength ) elif inpaint_image: image = inpaint_image['background'] if not all(pixel == (0, 0, 0) for pixel in list(inpaint_image['background'].getdata())) else None mask_image = inpaint_image['layers'][0] if image else None strength = float(inpaint_strength) if image and mask_image: base_args = BaseInpaintReq( **base_args.__dict__, image=image, mask_image=mask_image, strength=strength ) elif any([canny_image, pose_image, depth_image]): base_args.controlnet_config = ControlNetReq( controlnets=[], control_images=[], controlnet_conditioning_scale=[] ) if canny_image: base_args.controlnet_config.controlnets.append("canny") base_args.controlnet_config.control_images.append(canny_image) base_args.controlnet_config.controlnet_conditioning_scale.append(float(canny_strength)) if pose_image: base_args.controlnet_config.controlnets.append("pose") base_args.controlnet_config.control_images.append(pose_image) base_args.controlnet_config.controlnet_conditioning_scale.append(float(pose_strength)) if depth_image: base_args.controlnet_config.controlnets.append("depth") base_args.controlnet_config.control_images.append(depth_image) base_args.controlnet_config.controlnet_conditioning_scale.append(float(depth_strength)) if model in Config.IMAGES_MODELS and model['loader'] == 'sdxl' and scribble_image: base_args.controlnet_config.controlnets.append("scribble") base_args.controlnet_config.control_images.append(scribble_image) base_args.controlnet_config.controlnet_conditioning_scale.append(float(scribble_strength)) else: base_args = BaseReq(**base_args.__dict__) return gr.update( value=gen_img(base_args, progress), interactive=True ) except Exception as e: raise gr.Error(f"Error: {e}") from e