import os import gradio as gr import json import logging import torch from PIL import Image import spaces from diffusers import DiffusionPipeline, AutoencoderTiny, AutoencoderKL, AutoPipelineForImage2Image, FluxControlNetModel from diffusers.pipelines import FluxControlNetPipeline 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 huggingface_hub import hf_hub_download, HfFileSystem, ModelCard, snapshot_download import copy import random import time import requests import pandas as pd from transformers import pipeline from gradio_imageslider import ImageSlider import numpy as np import warnings huggingface_token = os.getenv("HUGGINFACE_TOKEN") translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en", device="cpu") #Load prompts for randomization df = pd.read_csv('prompts.csv', header=None) prompt_values = df.values.flatten() # Load LoRAs from JSON file with open('loras.json', 'r') as f: loras = json.load(f) # Initialize the base model dtype = torch.bfloat16 device = "cuda" if torch.cuda.is_available() else "cpu" # 공통 FLUX 모델 로드 base_model = "black-forest-labs/FLUX.1-dev" pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=dtype).to(device) # LoRA를 위한 설정 taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).to(device) good_vae = AutoencoderKL.from_pretrained(base_model, subfolder="vae", torch_dtype=dtype).to(device) # Image-to-Image 파이프라인 설정 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 ).to(device) # Upscale을 위한 ControlNet 설정 controlnet = FluxControlNetModel.from_pretrained( "jasperai/Flux.1-dev-Controlnet-Upscaler", torch_dtype=torch.bfloat16 ).to(device) # Upscale 파이프라인 설정 (기존 pipe 재사용) pipe_upscale = FluxControlNetPipeline( vae=pipe.vae, text_encoder=pipe.text_encoder, text_encoder_2=pipe.text_encoder_2, tokenizer=pipe.tokenizer, tokenizer_2=pipe.tokenizer_2, transformer=pipe.transformer, scheduler=pipe.scheduler, controlnet=controlnet ).to(device) MAX_SEED = 2**32 - 1 MAX_PIXEL_BUDGET = 1024 * 1024 pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe) class calculateDuration: def __init__(self, activity_name=""): self.activity_name = activity_name def __enter__(self): self.start_time = time.time() return self def __exit__(self, exc_type, exc_value, traceback): self.end_time = time.time() self.elapsed_time = self.end_time - self.start_time if self.activity_name: print(f"Elapsed time for {self.activity_name}: {self.elapsed_time:.6f} seconds") else: print(f"Elapsed time: {self.elapsed_time:.6f} seconds") def download_file(url, directory=None): if directory is None: directory = os.getcwd() # Use current working directory if not specified # Get the filename from the URL filename = url.split('/')[-1] # Full path for the downloaded file filepath = os.path.join(directory, filename) # Download the file response = requests.get(url) response.raise_for_status() # Raise an exception for bad status codes # Write the content to the file with open(filepath, 'wb') as file: file.write(response.content) return filepath def update_selection(evt: gr.SelectData, selected_indices, loras_state, width, height): selected_index = evt.index selected_indices = selected_indices or [] if selected_index in selected_indices: selected_indices.remove(selected_index) else: if len(selected_indices) < 2: selected_indices.append(selected_index) else: gr.Warning("You can select up to 2 LoRAs, remove one to select a new one.") return gr.update(), gr.update(), gr.update(), selected_indices, gr.update(), gr.update(), width, height, gr.update(), gr.update() selected_info_1 = "Select a LoRA 1" selected_info_2 = "Select a LoRA 2" lora_scale_1 = 1.15 lora_scale_2 = 1.15 lora_image_1 = None lora_image_2 = None if len(selected_indices) >= 1: lora1 = loras_state[selected_indices[0]] selected_info_1 = f"### LoRA 1 Selected: [{lora1['title']}](https://huggingface.co/{lora1['repo']}) ✨" lora_image_1 = lora1['image'] if len(selected_indices) >= 2: lora2 = loras_state[selected_indices[1]] selected_info_2 = f"### LoRA 2 Selected: [{lora2['title']}](https://huggingface.co/{lora2['repo']}) ✨" lora_image_2 = lora2['image'] if selected_indices: last_selected_lora = loras_state[selected_indices[-1]] new_placeholder = f"Type a prompt for {last_selected_lora['title']}" else: new_placeholder = "Type a prompt after selecting a LoRA" return gr.update(placeholder=new_placeholder), selected_info_1, selected_info_2, selected_indices, lora_scale_1, lora_scale_2, width, height, lora_image_1, lora_image_2 def remove_lora_1(selected_indices, loras_state): if len(selected_indices) >= 1: selected_indices.pop(0) selected_info_1 = "Select a LoRA 1" selected_info_2 = "Select a LoRA 2" lora_scale_1 = 1.15 lora_scale_2 = 1.15 lora_image_1 = None lora_image_2 = None if len(selected_indices) >= 1: lora1 = loras_state[selected_indices[0]] selected_info_1 = f"### LoRA 1 Selected: [{lora1['title']}]({lora1['repo']}) ✨" lora_image_1 = lora1['image'] if len(selected_indices) >= 2: lora2 = loras_state[selected_indices[1]] selected_info_2 = f"### LoRA 2 Selected: [{lora2['title']}]({lora2['repo']}) ✨" lora_image_2 = lora2['image'] return selected_info_1, selected_info_2, selected_indices, lora_scale_1, lora_scale_2, lora_image_1, lora_image_2 def remove_lora_2(selected_indices, loras_state): if len(selected_indices) >= 2: selected_indices.pop(1) selected_info_1 = "Select LoRA 1" selected_info_2 = "Select LoRA 2" lora_scale_1 = 1.15 lora_scale_2 = 1.15 lora_image_1 = None lora_image_2 = None if len(selected_indices) >= 1: lora1 = loras_state[selected_indices[0]] selected_info_1 = f"### LoRA 1 Selected: [{lora1['title']}]({lora1['repo']}) ✨" lora_image_1 = lora1['image'] if len(selected_indices) >= 2: lora2 = loras_state[selected_indices[1]] selected_info_2 = f"### LoRA 2 Selected: [{lora2['title']}]({lora2['repo']}) ✨" lora_image_2 = lora2['image'] return selected_info_1, selected_info_2, selected_indices, lora_scale_1, lora_scale_2, lora_image_1, lora_image_2 def randomize_loras(selected_indices, loras_state): try: if len(loras_state) < 2: raise gr.Error("Not enough LoRAs to randomize.") selected_indices = random.sample(range(len(loras_state)), 2) lora1 = loras_state[selected_indices[0]] lora2 = loras_state[selected_indices[1]] selected_info_1 = f"### LoRA 1 Selected: [{lora1['title']}](https://huggingface.co/{lora1['repo']}) ✨" selected_info_2 = f"### LoRA 2 Selected: [{lora2['title']}](https://huggingface.co/{lora2['repo']}) ✨" lora_scale_1 = 1.15 lora_scale_2 = 1.15 lora_image_1 = lora1['image'] lora_image_2 = lora2['image'] random_prompt = random.choice(prompt_values) return selected_info_1, selected_info_2, selected_indices, lora_scale_1, lora_scale_2, lora_image_1, lora_image_2, random_prompt except Exception as e: print(f"Error in randomize_loras: {str(e)}") return "Error", "Error", [], 1.15, 1.15, None, None, "" def add_custom_lora(custom_lora, selected_indices, current_loras): if custom_lora: try: title, repo, path, trigger_word, image = check_custom_model(custom_lora) 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: if repo.endswith(".safetensors") and repo.startswith("http"): repo = download_file(repo) new_item = { "image": image if image else "/home/user/app/custom.png", "title": title, "repo": repo, "weights": path, "trigger_word": trigger_word } print(f"New LoRA: {new_item}") existing_item_index = len(current_loras) current_loras.append(new_item) # Update gallery gallery_items = [(item["image"], item["title"]) for item in current_loras] # Update selected_indices if there's room if len(selected_indices) < 2: selected_indices.append(existing_item_index) else: gr.Warning("You can select up to 2 LoRAs, remove one to select a new one.") # Update selected_info and images selected_info_1 = "Select a LoRA 1" selected_info_2 = "Select a LoRA 2" lora_scale_1 = 1.15 lora_scale_2 = 1.15 lora_image_1 = None lora_image_2 = None if len(selected_indices) >= 1: lora1 = current_loras[selected_indices[0]] selected_info_1 = f"### LoRA 1 Selected: {lora1['title']} ✨" lora_image_1 = lora1['image'] if lora1['image'] else None if len(selected_indices) >= 2: lora2 = current_loras[selected_indices[1]] selected_info_2 = f"### LoRA 2 Selected: {lora2['title']} ✨" lora_image_2 = lora2['image'] if lora2['image'] else None print("Finished adding custom LoRA") return ( current_loras, gr.update(value=gallery_items), selected_info_1, selected_info_2, selected_indices, lora_scale_1, lora_scale_2, lora_image_1, lora_image_2 ) except Exception as e: print(e) gr.Warning(str(e)) return current_loras, gr.update(), gr.update(), gr.update(), selected_indices, gr.update(), gr.update(), gr.update(), gr.update() else: return current_loras, gr.update(), gr.update(), gr.update(), selected_indices, gr.update(), gr.update(), gr.update(), gr.update() def remove_custom_lora(selected_indices, current_loras): if current_loras: custom_lora_repo = current_loras[-1]['repo'] # Remove from loras list current_loras = current_loras[:-1] # Remove from selected_indices if selected custom_lora_index = len(current_loras) if custom_lora_index in selected_indices: selected_indices.remove(custom_lora_index) # Update gallery gallery_items = [(item["image"], item["title"]) for item in current_loras] # Update selected_info and images selected_info_1 = "Select a LoRA 1" selected_info_2 = "Select a LoRA 2" lora_scale_1 = 1.15 lora_scale_2 = 1.15 lora_image_1 = None lora_image_2 = None if len(selected_indices) >= 1: lora1 = current_loras[selected_indices[0]] selected_info_1 = f"### LoRA 1 Selected: [{lora1['title']}]({lora1['repo']}) ✨" lora_image_1 = lora1['image'] if len(selected_indices) >= 2: lora2 = current_loras[selected_indices[1]] selected_info_2 = f"### LoRA 2 Selected: [{lora2['title']}]({lora2['repo']}) ✨" lora_image_2 = lora2['image'] return ( current_loras, gr.update(value=gallery_items), selected_info_1, selected_info_2, selected_indices, lora_scale_1, lora_scale_2, lora_image_1, lora_image_2 ) @spaces.GPU(duration=75) def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, progress): print("Generating image...") pipe.to("cuda") generator = torch.Generator(device="cuda").manual_seed(seed) with calculateDuration("Generating image"): # Generate image for img in pipe.flux_pipe_call_that_returns_an_iterable_of_images( prompt=prompt_mash, num_inference_steps=steps, guidance_scale=cfg_scale, width=width, height=height, generator=generator, joint_attention_kwargs={"scale": 1.0}, output_type="pil", good_vae=good_vae, ): yield img @spaces.GPU(duration=75) def generate_image_to_image(prompt_mash, image_input_path, image_strength, steps, cfg_scale, width, height, seed): pipe_i2i.to("cuda") generator = torch.Generator(device="cuda").manual_seed(seed) image_input = load_image(image_input_path) final_image = pipe_i2i( prompt=prompt_mash, image=image_input, strength=image_strength, num_inference_steps=steps, guidance_scale=cfg_scale, width=width, height=height, generator=generator, joint_attention_kwargs={"scale": 1.0}, output_type="pil", ).images[0] return final_image def run_lora(prompt, image_input, image_strength, cfg_scale, steps, selected_indices, lora_scale_1, lora_scale_2, randomize_seed, seed, width, height, loras_state, progress=gr.Progress(track_tqdm=True)): try: # 한글 감지 및 번역 if any('\u3131' <= char <= '\u318E' or '\uAC00' <= char <= '\uD7A3' for char in prompt): translated = translator(prompt, max_length=512)[0]['translation_text'] print(f"Original prompt: {prompt}") print(f"Translated prompt: {translated}") prompt = translated if not selected_indices: raise gr.Error("You must select at least one LoRA before proceeding.") selected_loras = [loras_state[idx] for idx in selected_indices] # Build the prompt with trigger words prepends = [] appends = [] for lora in selected_loras: trigger_word = lora.get('trigger_word', '') if trigger_word: if lora.get("trigger_position") == "prepend": prepends.append(trigger_word) else: appends.append(trigger_word) prompt_mash = " ".join(prepends + [prompt] + appends) print("Prompt Mash: ", prompt_mash) # Unload previous LoRA weights with calculateDuration("Unloading LoRA"): pipe.unload_lora_weights() pipe_i2i.unload_lora_weights() print(pipe.get_active_adapters()) # Load LoRA weights with respective scales lora_names = [] lora_weights = [] with calculateDuration("Loading LoRA weights"): for idx, lora in enumerate(selected_loras): lora_name = f"lora_{idx}" lora_names.append(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) else: pipe_i2i.load_lora_weights(lora_path, low_cpu_mem_usage=True, adapter_name=lora_name) else: if weight_name: pipe.load_lora_weights(lora_path, weight_name=weight_name, low_cpu_mem_usage=True, adapter_name=lora_name) else: pipe.load_lora_weights(lora_path, low_cpu_mem_usage=True, adapter_name=lora_name) print("Loaded LoRAs:", lora_names) print("Adapter weights:", lora_weights) if image_input is not None: pipe_i2i.set_adapters(lora_names, adapter_weights=lora_weights) else: pipe.set_adapters(lora_names, adapter_weights=lora_weights) print(pipe.get_active_adapters()) # Set random seed for reproducibility with calculateDuration("Randomizing seed"): if randomize_seed: seed = random.randint(0, MAX_SEED) # Generate image if image_input is not None: final_image = generate_image_to_image(prompt_mash, image_input, image_strength, steps, cfg_scale, width, height, seed) else: image_generator = generate_image(prompt_mash, steps, seed, cfg_scale, width, height, progress) final_image = None step_counter = 0 for image in image_generator: step_counter += 1 final_image = image progress_bar = f'