SemaSci/DiffModels

1# не запускается, ошибка с кешем, отложил пока2 3# это файл только с LoRA, без ControlNet и IpAdapter4 5import gradio as gr6import numpy as np7import random8 9# import spaces #[uncomment to use ZeroGPU]10from diffusers import DiffusionPipeline11import torch12 13from peft import PeftModel, LoraConfig14import os15 16# Добавляем глобальный кэш для пайплайнов17pipe_cache = {}18 19def get_lora_sd_pipeline(20    ckpt_dir='./lora_logos', 21    base_model_name_or_path=None, 22    dtype=torch.float16, 23    adapter_name="default"24    ):25    26    unet_sub_dir = os.path.join(ckpt_dir, "unet")27    text_encoder_sub_dir = os.path.join(ckpt_dir, "text_encoder")28    29    if os.path.exists(text_encoder_sub_dir) and base_model_name_or_path is None:30        config = LoraConfig.from_pretrained(text_encoder_sub_dir)31        base_model_name_or_path = config.base_model_name_or_path32    33    if base_model_name_or_path is None:34        raise ValueError("Please specify the base model name or path")35    36    pipe = DiffusionPipeline.from_pretrained(base_model_name_or_path, torch_dtype=dtype)37    before_params = pipe.unet.parameters()38#    pipe.unet = PeftModel.from_pretrained(pipe.unet, unet_sub_dir, adapter_name=adapter_name)39    # Исправляем загрузку конфигурации40    config = LoraConfig.from_pretrained(unet_sub_dir)41    42    pipe.unet = PeftModel.from_pretrained(43        pipe.unet, 44        unet_sub_dir, 45        adapter_name=adapter_name,46        config=config  # Явно передаем конфигурацию47    )48    49    pipe.unet.set_adapter(adapter_name)50    after_params = pipe.unet.parameters()51    print("UNet Parameters changed:", any(torch.any(b != a) for b, a in zip(before_params, after_params)))52    53    if os.path.exists(text_encoder_sub_dir):54        pipe.text_encoder = PeftModel.from_pretrained(pipe.text_encoder, text_encoder_sub_dir, adapter_name=adapter_name)55    56    if dtype in (torch.float16, torch.bfloat16):57        pipe.unet.half()58        if pipe.text_encoder is not None:59            pipe.text_encoder.half()60    61    return pipe62 63def process_prompt(prompt, tokenizer, text_encoder, max_length=77):64    tokens = tokenizer(prompt, truncation=False, return_tensors="pt")["input_ids"]65    chunks = [tokens[:, i:i + max_length] for i in range(0, tokens.shape[1], max_length)]66    67    with torch.no_grad():68        embeds = [text_encoder(chunk.to(text_encoder.device))[0] for chunk in chunks]69    70    return torch.cat(embeds, dim=1)71 72def align_embeddings(prompt_embeds, negative_prompt_embeds):73    max_length = max(prompt_embeds.shape[1], negative_prompt_embeds.shape[1])74    return torch.nn.functional.pad(prompt_embeds, (0, 0, 0, max_length - prompt_embeds.shape[1])), \75           torch.nn.functional.pad(negative_prompt_embeds, (0, 0, 0, max_length - negative_prompt_embeds.shape[1]))76 77device = "cuda" if torch.cuda.is_available() else "cpu"78#model_repo_id = "stabilityai/sdxl-turbo"  # Replace to the model you would like to use79model_id_default = "sd-legacy/stable-diffusion-v1-5"80model_dropdown = ['stabilityai/sdxl-turbo', 'CompVis/stable-diffusion-v1-4', 'sd-legacy/stable-diffusion-v1-5'  ]81 82model_lora_default = "lora_pussinboots_logos"83model_lora_dropdown = ['lora_lady_and_cats_logos', 'lora_pussinboots_logos'  ]84 85if torch.cuda.is_available():86    torch_dtype = torch.float1687else:88    torch_dtype = torch.float3289 90# pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)91# pipe = pipe.to(device)92 93MAX_SEED = np.iinfo(np.int32).max94MAX_IMAGE_SIZE = 102495 96 97# @spaces.GPU #[uncomment to use ZeroGPU]98def infer(99    prompt,100    negative_prompt,101    randomize_seed,102    width=512,103    height=512,104    model_repo_id=model_id_default,105    seed=42,106    guidance_scale=7,107    num_inference_steps=20,108    model_lora_id=model_lora_default,109    lora_scale=0.5,110    progress=gr.Progress(track_tqdm=True),111    ):112        113    global pipe_cache114 115    if randomize_seed:116        seed = random.randint(0, MAX_SEED)117 118    generator = torch.Generator().manual_seed(seed)119 120    # Кэширование пайплайнов121    cache_key = f"{model_repo_id}_{model_lora_id}"122    if cache_key not in pipe_cache:123        if model_repo_id != model_id_default:124            pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype).to(device)125            prompt_embeds = process_prompt(prompt, pipe.tokenizer, pipe.text_encoder)126            negative_prompt_embeds = process_prompt(negative_prompt, pipe.tokenizer, pipe.text_encoder)127            prompt_embeds, negative_prompt_embeds = align_embeddings(prompt_embeds, negative_prompt_embeds)128        else:129            pipe = get_lora_sd_pipeline(130                ckpt_dir='./'+model_lora_id,131                base_model_name_or_path=model_id_default,132                dtype=torch_dtype133            ).to(device)134        135        pipe_cache[cache_key] = pipe136    else:137        pipe = pipe_cache[cache_key]138 139    # Динамическое применение масштаба LoRA140    if model_repo_id == model_id_default:141        # Убираем fuse_lora()142        # pipe.fuse_lora(lora_scale=lora_scale)  # Закомментировали проблемную строку143        144        # Вместо этого устанавливаем адаптеры динамически145        pipe.unet.set_adapters(146            [model_lora_id], 147            adapter_weights=[lora_scale]148        )149        if hasattr(pipe, 'text_encoder') and pipe.text_encoder is not None:150            pipe.text_encoder.set_adapters(151                [model_lora_id], 152                adapter_weights=[lora_scale]153            )154        155        print(f"Active adapters - UNet: {pipe.unet.active_adapters}, Text Encoder: {pipe.text_encoder.active_adapters if hasattr(pipe, 'text_encoder') else None}")156        print("UNet first layer weights:", pipe.unet.base_model.model[0].weight.data[0,0,:5])157        print(f"LoRA scale applied: {lora_scale}")158    159    160    # на вызов pipe с эмбеддингами161    params = {162        'prompt_embeds': prompt_embeds,163        'negative_prompt_embeds': negative_prompt_embeds,164        'guidance_scale': guidance_scale,165        'num_inference_steps': num_inference_steps,166        'width': width,167        'height': height,168        'generator': generator,169    }170    171    return pipe(**params).images[0], seed    172 173    # return image, seed174 175 176examples = [177    "Puss in Boots wearing a sombrero crosses the Grand Canyon on a tightrope with a guitar.",178    "A cat is playing a song called ""About the Cat"" on an accordion by the sea at sunset. The sun is quickly setting behind the horizon, and the light is fading.",179    "A cat walks through the grass on the streets of an abandoned city. The camera view is always focused on the cat's face.",180    "A young lady in a Russian embroidered kaftan is sitting on a beautiful carved veranda, holding a cup to her mouth and drinking tea from the cup. With her other hand, the girl holds a saucer. The cup and saucer are painted with gzhel. Next to the girl on the table stands a samovar, and steam can be seen above it.",181    "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",182    "An astronaut riding a green horse",183    "A delicious ceviche cheesecake slice",184]185 186css = """187#col-container {188    margin: 0 auto;189    max-width: 640px;190}191"""192 193with gr.Blocks(css=css) as demo:194    with gr.Column(elem_id="col-container"):195        gr.Markdown(" # Text-to-Image SemaSci Template")196 197        with gr.Row():198            prompt = gr.Text(199                label="Prompt",200                show_label=False,201                max_lines=1,202                placeholder="Enter your prompt",203                container=False,204            )205 206            run_button = gr.Button("Run", scale=0, variant="primary")207 208        result = gr.Image(label="Result", show_label=False)209 210        with gr.Accordion("Advanced Settings", open=False):211#            model_repo_id = gr.Text(212#                label="Model Id",213#                max_lines=1,214#                placeholder="Choose model",215#                visible=True,216#                value=model_repo_id,217#            )218            model_repo_id = gr.Dropdown(219                label="Model Id",220                choices=model_dropdown,221                info="Choose model",222                visible=True,223                allow_custom_value=True,224#                value=model_repo_id,225                value=model_id_default,226            )            227            228            negative_prompt = gr.Text(229                label="Negative prompt",230                max_lines=1,231                placeholder="Enter a negative prompt",232                visible=True,233            )234 235            seed = gr.Slider(236                label="Seed",237                minimum=0,238                maximum=MAX_SEED,239                step=1,240                value=42,241            )242 243            randomize_seed = gr.Checkbox(label="Randomize seed", value=False)244 245            with gr.Row():246                width = gr.Slider(247                    label="Width",248                    minimum=256,249                    maximum=MAX_IMAGE_SIZE,250                    step=32,251                    value=256,  # Replace with defaults that work for your model252                )253 254                height = gr.Slider(255                    label="Height",256                    minimum=256,257                    maximum=MAX_IMAGE_SIZE,258                    step=32,259                    value=256,  # Replace with defaults that work for your model260                )261 262            with gr.Row():263                guidance_scale = gr.Slider(264                    label="Guidance scale",265                    minimum=0.0,266                    maximum=10.0,267                    step=0.1,268                    value=7.0,  # Replace with defaults that work for your model269                )270 271                num_inference_steps = gr.Slider(272                    label="Number of inference steps",273                    minimum=1,274                    maximum=50,275                    step=1,276                    value=20,  # Replace with defaults that work for your model277                )278                279            with gr.Row():280                model_lora_id = gr.Dropdown(281                    label="Lora Id",282                    choices=model_lora_dropdown,283                    info="Choose LoRA model",284                    visible=True,285                    allow_custom_value=True,286                    value=model_lora_default,287                )            288            289                lora_scale = gr.Slider(290                    label="LoRA scale",291                    minimum=0.0,292                    maximum=1.0,293                    step=0.1,294                    value=0.5,295                )                296 297        gr.Examples(examples=examples, inputs=[prompt])298    gr.on(299        triggers=[run_button.click, prompt.submit],300        fn=infer,301        inputs=[302            prompt,303            negative_prompt,304            randomize_seed,305            width,306            height,307            model_repo_id,308            seed,309            guidance_scale,310            num_inference_steps,311            model_lora_id,312            lora_scale,313        ],314        outputs=[result, seed],315    )316 317if __name__ == "__main__":318    demo.launch()