import gradio as gr import numpy as np import random # import spaces #[uncomment to use ZeroGPU] from diffusers import DiffusionPipeline from peft import PeftModel, PeftConfig from rembg import remove from PIL import Image import io import torch from typing import Optional # кэш для пайплайнов (чтобы не перезагружать модель при каждом запросе) PIPE_CACHE: dict[str, DiffusionPipeline] = {} DEFAULT_MODEL = "CompVis/stable-diffusion-v1-4" BASE_MODEL_FOR_LORA = "stable-diffusion-v1-5/stable-diffusion-v1-5" # Base model used for LoRA training LORA_MODEL_ID = "DiZH797/my-tuned-lora" # Your uploaded LoRA model ID MODEL_OPTIONS = [ "CompVis/stable-diffusion-v1-4", "stabilityai/stable-diffusion-2-1", "stabilityai/sdxl-turbo", LORA_MODEL_ID ] device = "cuda" if torch.cuda.is_available() else "cpu" model_repo_id = "stabilityai/sdxl-turbo" # Replace to the model you would like to use if torch.cuda.is_available(): torch_dtype = torch.float16 else: torch_dtype = torch.float32 # pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype) # pipe = pipe.to(device) MAX_SEED = np.iinfo(np.int32).max MAX_IMAGE_SIZE = 1024 def get_pipe(model_id: str, lora_scale: float = 1.0): """ Loads the pipeline for a given model ID. If the selected model is the LoRA, it loads the base model and then merges the LoRA weights. """ cache_key = f"{model_id}_{lora_scale}" if cache_key in PIPE_CACHE: return PIPE_CACHE[cache_key] # Check if the selected model is the LoRA adapter if model_id == LORA_MODEL_ID: # Укажите правильные имена файлов pipe = DiffusionPipeline.from_pretrained( BASE_MODEL_FOR_LORA, dtype=torch_dtype ).to(device) # pipe.unet = PeftModel.from_pretrained(pipe.unet, LORA_MODEL_ID) pipe.load_lora_weights( LORA_MODEL_ID, weight_name="merged_lora_weights.safetensors" ) pipe.fuse_lora(lora_scale=lora_scale) # После загрузки LoRA print("LoRa scale is", lora_scale) print("LoRA layers in unet:") for name, param in pipe.unet.named_parameters(): if "lora" in name.lower() and param.requires_grad: print(f"Unet LoRA layer: {name}, shape: {param.shape}") break print("LoRA layers in text_encoder:") for name, param in pipe.text_encoder.named_parameters(): if "lora" in name: print(f"Text Encoder LoRA: {name}, shape: {param.shape}") break else: # Load a standard model without LoRA pipe = DiffusionPipeline.from_pretrained( model_id, dtype=torch_dtype ).to(device) PIPE_CACHE[cache_key] = pipe return pipe # @spaces.GPU #[uncomment to use ZeroGPU] def infer( model_id: Optional[str] = DEFAULT_MODEL, prompt: str = "", negative_prompt: str = "", seed: int = 42, randomize_seed: bool = False, width: int = 512, height: int = 512, guidance_scale: float = 7.0, num_inference_steps: int = 20, scheduler_name: Optional[str] = None, lora_scale: float = 1.0, remove_background: bool = False, progress=gr.Progress(track_tqdm=True), ): # получаем/загружаем нужный pipe pipe = get_pipe(model_id, lora_scale) # при желании можно подменить scheduler по имени (опционально) if scheduler_name: # примерная схема: словарь name->класс scheduler # при необходимости добавить другие scheduler'ы — импортируйте их сверху и добавьте сюда try: from diffusers import DDIMScheduler, EulerAncestralDiscreteScheduler, PNDMScheduler, DPMSolverMultistepScheduler sched_map = { "DDIM": DDIMScheduler, "EulerAncestral": EulerAncestralDiscreteScheduler, "PNDM": PNDMScheduler, "DPMSMS": DPMSolverMultistepScheduler } if scheduler_name in sched_map: pipe.scheduler = sched_map[scheduler_name].from_config(pipe.scheduler.config) except Exception: # если что-то пошло не так — просто используем дефолтный scheduler pass if randomize_seed: seed = random.randint(0, MAX_SEED) generator = torch.Generator().manual_seed(int(seed)) image = pipe( prompt=prompt, negative_prompt=negative_prompt, guidance_scale=guidance_scale, num_inference_steps=num_inference_steps, width=width, height=height, generator=generator, ).images[0] if remove_background: # Конвертируем PIL Image в bytes img_byte_arr = io.BytesIO() image.save(img_byte_arr, format='PNG') img_byte_arr = img_byte_arr.getvalue() # Удаляем фон output_image = remove(img_byte_arr) # Конвертируем обратно в PIL Image image = Image.open(io.BytesIO(output_image)) return image, seed examples = [ "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k", "An astronaut riding a green horse", "A delicious ceviche cheesecake slice", ] css = """ #col-container { margin: 0 auto; max-width: 640px; } """ with gr.Blocks(css=css) as demo: with gr.Column(elem_id="col-container"): gr.Markdown(" # Text-to-Image Gradio Template") # Model selector (выпадающий список) model_select = gr.Dropdown( label="Model", choices=MODEL_OPTIONS, value=DEFAULT_MODEL, interactive=True, ) # опциональный селектор scheduler scheduler_select = gr.Dropdown( label="Scheduler (optional)", choices=["", "DDIM", "EulerAncestral", "PNDM", "DPMSMS"], value="", ) # Add a new slider for LoRA scale lora_scale_slider = gr.Slider( label="LoRA Scale (Only for LoRA model)", minimum=0.0, maximum=3.0, step=0.1, value=1.0, visible=False, # Initially hidden ) with gr.Row(): prompt = gr.Text( label="Prompt", show_label=False, max_lines=1, placeholder="Enter your prompt", container=False, ) run_button = gr.Button("Run", scale=0, variant="primary") result = gr.Image(label="Result", show_label=False) with gr.Accordion("Advanced Settings", open=False): remove_background = gr.Checkbox( label="Remove background from generated image", value=False, info="Use rembg to remove background from the generated image" ) negative_prompt = gr.Text( label="Negative prompt", max_lines=1, placeholder="Enter a negative prompt", visible=True, ) seed = gr.Slider( label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42, ) randomize_seed = gr.Checkbox(label="Randomize seed", value=True) with gr.Row(): width = gr.Slider( label="Width", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=512, # Replace with defaults that work for your model ) height = gr.Slider( label="Height", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=512, # Replace with defaults that work for your model ) with gr.Row(): guidance_scale = gr.Slider( label="Guidance scale", minimum=0.0, maximum=10.0, step=0.1, value=7.0, # Replace with defaults that work for your model ) num_inference_steps = gr.Slider( label="Number of inference steps", minimum=1, maximum=50, step=1, value=20, # Replace with defaults that work for your model ) gr.Examples(examples=examples, inputs=[prompt]) # Function to show/hide the LoRA scale slider based on model selection def toggle_lora_scale_slider(model_id): if model_id == LORA_MODEL_ID: return gr.Slider(visible=True) else: return gr.Slider(visible=False) model_select.change( fn=toggle_lora_scale_slider, inputs=model_select, outputs=lora_scale_slider ) gr.on( triggers=[run_button.click, prompt.submit], fn=infer, inputs=[ model_select, prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, scheduler_select, lora_scale_slider, remove_background ], outputs=[result, seed], ) if __name__ == "__main__": demo.launch()