Upload app.py

app.py

@@ -5,7 +5,7 @@ import torch
 from diffusers import DiffusionPipeline, AutoencoderTiny, AutoencoderKL, AutoPipelineForImage2Image
 from live_preview_helpers import calculate_shift, retrieve_timesteps, flux_pipe_call_that_returns_an_iterable_of_images
 from diffusers.utils import load_image
-from diffusers import FluxControlNetPipeline, FluxControlNetModel, FluxMultiControlNetModel, FluxControlNetImg2ImgPipeline
+from diffusers import FluxControlNetPipeline, FluxControlNetModel, FluxMultiControlNetModel, FluxControlNetImg2ImgPipeline, FluxTransformer2DModel
 from huggingface_hub import hf_hub_download, HfFileSystem, ModelCard, snapshot_download, HfApi
 import os
 import copy
@@ -13,8 +13,9 @@ import random
 import time
 import requests
 import pandas as pd
+from pathlib import Path
 
-from env import models, num_loras, num_cns, HF_TOKEN
+from env import models, num_loras, num_cns, HF_TOKEN, single_file_base_models
 from mod import (clear_cache, get_repo_safetensors, is_repo_name, is_repo_exists, get_model_trigger,
                  description_ui, compose_lora_json, is_valid_lora, fuse_loras, save_image, preprocess_i2i_image,
                  get_trigger_word, enhance_prompt, set_control_union_image,
@@ -41,11 +42,11 @@ controlnet_model_union_repo = 'InstantX/FLUX.1-dev-Controlnet-Union'
 dtype = torch.bfloat16
 #dtype = torch.float8_e4m3fn
 #device = "cuda" if torch.cuda.is_available() else "cpu"
-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)
+taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype, token=HF_TOKEN)
+good_vae = AutoencoderKL.from_pretrained(base_model, subfolder="vae", torch_dtype=dtype, token=HF_TOKEN)
+pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=dtype, vae=taef1, token=HF_TOKEN)
 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)
+ tokenizer=pipe.tokenizer, text_encoder_2=pipe.text_encoder_2, tokenizer_2=pipe.tokenizer_2, torch_dtype=dtype, token=HF_TOKEN)
 controlnet_union = None
 controlnet = None
 last_model = models[0]
@@ -70,10 +71,13 @@ def unload_lora():
 # https://huggingface.co/spaces/jiuface/FLUX.1-dev-Controlnet-Union
 # https://huggingface.co/docs/diffusers/main/en/api/pipelines/flux
 #@spaces.GPU()
-def change_base_model(repo_id: str, cn_on: bool, disable_model_cache: bool, progress=gr.Progress(track_tqdm=True)):
+def change_base_model(repo_id: str, cn_on: bool, disable_model_cache: bool, model_type: str, progress=gr.Progress(track_tqdm=True)):
     global pipe, pipe_i2i, taef1, good_vae, controlnet_union, controlnet, last_model, last_cn_on, dtype
+    safetensors_file = None
+    single_file_base_model = single_file_base_models.get(model_type, models[0])
     try:
-        if not disable_model_cache and (repo_id == last_model and cn_on is last_cn_on) or not is_repo_name(repo_id) or not is_repo_exists(repo_id): return gr.update(visible=True)
+        #if not disable_model_cache and (repo_id == last_model and cn_on is last_cn_on) or not is_repo_name(repo_id) or not is_repo_exists(repo_id): return gr.update(visible=True)
+        if not disable_model_cache and (repo_id == last_model and cn_on is last_cn_on) or ((not is_repo_name(repo_id) or not is_repo_exists(repo_id)) and not ".safetensors" in repo_id): return gr.update(visible=True)
         unload_lora()
         pipe.to("cpu")
         pipe_i2i.to("cpu")
@@ -85,12 +89,19 @@ def change_base_model(repo_id: str, cn_on: bool, disable_model_cache: bool, prog
         if cn_on:
             progress(0, desc=f"Loading model: {repo_id} / Loading ControlNet: {controlnet_model_union_repo}")
             print(f"Loading model: {repo_id} / Loading ControlNet: {controlnet_model_union_repo}")
-            controlnet_union = FluxControlNetModel.from_pretrained(controlnet_model_union_repo, torch_dtype=dtype)
+            controlnet_union = FluxControlNetModel.from_pretrained(controlnet_model_union_repo, torch_dtype=dtype, token=HF_TOKEN)
             controlnet = FluxMultiControlNetModel([controlnet_union])
             controlnet.config = controlnet_union.config
-            pipe = FluxControlNetPipeline.from_pretrained(repo_id, controlnet=controlnet, torch_dtype=dtype)
-            pipe_i2i = FluxControlNetImg2ImgPipeline.from_pretrained(repo_id, controlnet=controlnet, vae=None, 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)
+            if ".safetensors" in repo_id:
+                safetensors_file = download_file_mod(repo_id)
+                transformer = FluxTransformer2DModel.from_single_file(safetensors_file, subfolder="transformer", torch_dtype=dtype, config=single_file_base_model)
+                pipe = FluxControlNetPipeline.from_pretrained(single_file_base_model, transformer=transformer, controlnet=controlnet, torch_dtype=dtype, token=HF_TOKEN)
+                pipe_i2i = FluxControlNetImg2ImgPipeline.from_pretrained(single_file_base_model, controlnet=controlnet, vae=None, 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, token=HF_TOKEN)
+            else:
+                pipe = FluxControlNetPipeline.from_pretrained(repo_id, controlnet=controlnet, torch_dtype=dtype, token=HF_TOKEN)
+                pipe_i2i = FluxControlNetImg2ImgPipeline.from_pretrained(repo_id, controlnet=controlnet, vae=None, 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, token=HF_TOKEN)
             last_model = repo_id
             last_cn_on = cn_on
             progress(1, desc=f"Model loaded: {repo_id} / ControlNet Loaded: {controlnet_model_union_repo}")
@@ -98,16 +109,25 @@ def change_base_model(repo_id: str, cn_on: bool, disable_model_cache: bool, prog
         else:
             progress(0, desc=f"Loading model: {repo_id}")
             print(f"Loading model: {repo_id}")
-            pipe = DiffusionPipeline.from_pretrained(repo_id, torch_dtype=dtype)
-            pipe_i2i = AutoPipelineForImage2Image.from_pretrained(repo_id, vae=None, 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)
+            if ".safetensors" in repo_id:
+                safetensors_file = download_file_mod(repo_id)
+                transformer = FluxTransformer2DModel.from_single_file(safetensors_file, subfolder="transformer", torch_dtype=dtype, config=single_file_base_model)
+                pipe = DiffusionPipeline.from_pretrained(single_file_base_model, transformer=transformer, torch_dtype=dtype, token=HF_TOKEN)
+                pipe_i2i = AutoPipelineForImage2Image.from_pretrained(single_file_base_model, vae=None, 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, token=HF_TOKEN)
+            else:
+                pipe = DiffusionPipeline.from_pretrained(repo_id, torch_dtype=dtype, token=HF_TOKEN)
+                pipe_i2i = AutoPipelineForImage2Image.from_pretrained(repo_id, vae=None, 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, token=HF_TOKEN)
             last_model = repo_id
             last_cn_on = cn_on
             progress(1, desc=f"Model loaded: {repo_id}")
             print(f"Model loaded: {repo_id}")
     except Exception as e:
-        print(f"Model load Error: {e}")
-        raise gr.Error(f"Model load Error: {e}") from e
+        print(f"Model load Error: {repo_id} {e}")
+        raise gr.Error(f"Model load Error: {repo_id} {e}") from e
+    finally:
+        if safetensors_file and Path(safetensors_file).exists(): Path(safetensors_file).unlink()
     return gr.update(visible=True)
 
 change_base_model.zerogpu = True
@@ -260,7 +280,6 @@ def add_custom_lora(custom_lora, selected_indices, current_loras, gallery):
                 except Exception as e:
                     print(e)
                     image = None
-            #image = image if is_repo_public(repo) else None
             print(f"Loaded custom LoRA: {repo}")
             existing_item_index = next((index for (index, item) in enumerate(current_loras) if item['repo'] == repo), None)
             if existing_item_index is None:
@@ -521,20 +540,27 @@ def run_lora(prompt, image_input, image_strength, cfg_scale, steps, selected_ind
         for idx, lora in enumerate(selected_loras):
             lora_name = f"lora_{idx}"
             lora_names.append(lora_name)
+            print(f"Lora Name: {lora_name}")
             lora_weights.append(lora_scale_1 if idx == 0 else lora_scale_2)
             lora_path = lora['repo']
             weight_name = lora.get("weights")
             print(f"Lora Path: {lora_path}")
             if image_input is not None:
-                if weight_name:
-                    pipe_i2i.load_lora_weights(lora_path, weight_name=weight_name, low_cpu_mem_usage=True, adapter_name=lora_name, token=HF_TOKEN)
-                else:
-                    pipe_i2i.load_lora_weights(lora_path, low_cpu_mem_usage=True, adapter_name=lora_name, token=HF_TOKEN)
+                pipe_i2i.load_lora_weights(
+                    lora_path, 
+                    weight_name=weight_name if weight_name else None,
+                    low_cpu_mem_usage=True,
+                    adapter_name=lora_name,
+                    token=HF_TOKEN
+                )
             else:
-                if weight_name:
-                    pipe.load_lora_weights(lora_path, weight_name=weight_name, low_cpu_mem_usage=True, adapter_name=lora_name, token=HF_TOKEN)
-                else:
-                    pipe.load_lora_weights(lora_path, low_cpu_mem_usage=True, adapter_name=lora_name, token=HF_TOKEN)
+                pipe.load_lora_weights(
+                    lora_path, 
+                    weight_name=weight_name if weight_name else None,
+                    low_cpu_mem_usage=True,
+                    adapter_name=lora_name,
+                    token=HF_TOKEN
+                )
         print("Loaded LoRAs:", lora_names)
         if image_input is not None:
             pipe_i2i.set_adapters(lora_names, adapter_weights=lora_weights)
@@ -722,8 +748,13 @@ with gr.Blocks(theme='NoCrypt/miku@>=1.2.2', fill_width=True, css=css, delete_ca
                 with gr.Accordion("History", open=False):
                     history_gallery = gr.Gallery(label="History", columns=6, object_fit="contain", interactive=False)
         with gr.Group():
-            model_name = gr.Dropdown(label="Base Model", info="You can enter a huggingface model repo_id to want to use.", choices=models, value=models[0], allow_custom_value=True)
+            with gr.Row():
+                model_name = gr.Dropdown(label="Base Model", info="You can enter a huggingface model repo_id or path of single safetensors file to want to use.",
+                                         choices=models, value=models[0], allow_custom_value=True, min_width=320, scale=5)
+                model_type = gr.Radio(label="Model type", info="Model type of single safetensors file",
+                                      choices=list(single_file_base_models.keys()), value=list(single_file_base_models.keys())[0], scale=1)
             model_info = gr.Markdown(elem_classes="info")
+            
         with gr.Row():
             with gr.Accordion("Advanced Settings", open=False):
                 with gr.Row():
@@ -756,7 +787,7 @@ with gr.Blocks(theme='NoCrypt/miku@>=1.2.2', fill_width=True, css=css, delete_ca
                             with gr.Row():
                                 for i in range(num_loras):
                                     with gr.Column():
-                                        lora_repo[i] = gr.Dropdown(label=f"LoRA {int(i+1)} Repo", choices=get_all_lora_tupled_list(), info="Input LoRA Repo ID", value="", allow_custom_value=True)
+                                        lora_repo[i] = gr.Dropdown(label=f"LoRA {int(i+1)} Repo", choices=get_all_lora_tupled_list(), info="Input LoRA Repo ID", value="", allow_custom_value=True, min_width=320)
                                         with gr.Row():
                                             lora_weights[i] = gr.Dropdown(label=f"LoRA {int(i+1)} Filename", choices=[], info="Optional", value="", allow_custom_value=True)
                                             lora_trigger[i] = gr.Textbox(label=f"LoRA {int(i+1)} Trigger Prompt", lines=1, max_lines=4, value="")
@@ -837,7 +868,7 @@ with gr.Blocks(theme='NoCrypt/miku@>=1.2.2', fill_width=True, css=css, delete_ca
     gr.on(
         triggers=[generate_button.click, prompt.submit],
         fn=change_base_model,
-        inputs=[model_name, cn_on, disable_model_cache],
+        inputs=[model_name, cn_on, disable_model_cache, model_type],
         outputs=[result],
         queue=True,
         show_api=False,