import os import gradio as gr import huggingface_hub import pillow_avif import spaces import torch import gc from huggingface_hub import snapshot_download from pillow_heif import register_heif_opener from pipelines.pipeline_infu_flux import InfUFluxPipeline # Register HEIF support for Pillow register_heif_opener() class ModelVersion: STAGE_1 = "sim_stage1" STAGE_2 = "aes_stage2" DEFAULT_VERSION = STAGE_2 ENABLE_ANTI_BLUR_DEFAULT = False ENABLE_REALISM_DEFAULT = False loaded_pipeline_config = { "model_version": "aes_stage2", "enable_realism": False, "enable_anti_blur": False, 'pipeline': None } def download_models(): snapshot_download(repo_id='ByteDance/InfiniteYou', local_dir='./models/InfiniteYou', local_dir_use_symlinks=False) try: snapshot_download(repo_id='black-forest-labs/FLUX.1-dev', local_dir='./models/FLUX.1-dev', local_dir_use_symlinks=False) except Exception as e: print(e) print('\nYou are downloading `black-forest-labs/FLUX.1-dev` to `./models/FLUX.1-dev` but failed. ' 'Please accept the agreement and obtain access at https://huggingface.co/black-forest-labs/FLUX.1-dev. ' 'Then, use `huggingface-cli login` and your access tokens at https://huggingface.co/settings/tokens to authenticate. ' 'After that, run the code again.') print('\nYou can also download it manually from HuggingFace and put it in `./models/InfiniteYou`, ' 'or you can modify `base_model_path` in `app.py` to specify the correct path.') exit() def init_pipeline(model_version, enable_realism, enable_anti_blur): loaded_pipeline_config["enable_realism"] = enable_realism loaded_pipeline_config["enable_anti_blur"] = enable_anti_blur loaded_pipeline_config["model_version"] = model_version pipeline = loaded_pipeline_config['pipeline'] gc.collect() torch.cuda.empty_cache() model_path = f'./models/InfiniteYou/infu_flux_v1.0/{model_version}' print(f'loading model from {model_path}') pipeline = InfUFluxPipeline( base_model_path='./models/FLUX.1-dev', infu_model_path=model_path, insightface_root_path='./models/InfiniteYou/supports/insightface', image_proj_num_tokens=8, infu_flux_version='v1.0', model_version=model_version, ) loaded_pipeline_config['pipeline'] = pipeline pipeline.pipe.delete_adapters(['realism', 'anti_blur']) loras = [] if enable_realism: loras.append(['realism', 1.0]) if enable_anti_blur: loras.append(['anti_blur', 1.0]) pipeline.load_loras_state_dict(loras) return pipeline def prepare_pipeline(model_version, enable_realism, enable_anti_blur): if ( loaded_pipeline_config['pipeline'] is not None and loaded_pipeline_config["enable_realism"] == enable_realism and loaded_pipeline_config["enable_anti_blur"] == enable_anti_blur and model_version == loaded_pipeline_config["model_version"] ): return loaded_pipeline_config['pipeline'] loaded_pipeline_config["enable_realism"] = enable_realism loaded_pipeline_config["enable_anti_blur"] = enable_anti_blur loaded_pipeline_config["model_version"] = model_version pipeline = loaded_pipeline_config['pipeline'] if pipeline is None or pipeline.model_version != model_version: print(f'Switching model to {model_version}') pipeline.model_version = model_version if model_version == 'aes_stage2': pipeline.infusenet_sim.cpu() pipeline.image_proj_model_sim.cpu() torch.cuda.empty_cache() pipeline.infusenet_aes.to('cuda') pipeline.pipe.controlnet = pipeline.infusenet_aes pipeline.image_proj_model_aes.to('cuda') pipeline.image_proj_model = pipeline.image_proj_model_aes else: pipeline.infusenet_aes.cpu() pipeline.image_proj_model_aes.cpu() torch.cuda.empty_cache() pipeline.infusenet_sim.to('cuda') pipeline.pipe.controlnet = pipeline.infusenet_sim pipeline.image_proj_model_sim.to('cuda') pipeline.image_proj_model = pipeline.image_proj_model_sim loaded_pipeline_config['pipeline'] = pipeline pipeline.pipe.delete_adapters(['realism', 'anti_blur']) loras = [] if enable_realism: loras.append(['realism', 1.0]) if enable_anti_blur: loras.append(['anti_blur', 1.0]) pipeline.load_loras_state_dict(loras) return pipeline @spaces.GPU(duration=120) def generate_image( input_image, control_image, prompt, seed, width, height, guidance_scale, num_steps, infusenet_conditioning_scale, infusenet_guidance_start, infusenet_guidance_end, enable_realism, enable_anti_blur, model_version ): try: pipeline = prepare_pipeline(model_version=model_version, enable_realism=enable_realism, enable_anti_blur=enable_anti_blur) if seed == 0: seed = torch.seed() & 0xFFFFFFFF image = pipeline( id_image=input_image, prompt=prompt, control_image=control_image, seed=seed, width=width, height=height, guidance_scale=guidance_scale, num_steps=num_steps, infusenet_conditioning_scale=infusenet_conditioning_scale, infusenet_guidance_start=infusenet_guidance_start, infusenet_guidance_end=infusenet_guidance_end, ) except Exception as e: print(e) gr.Error(f"An error occurred: {e}") return gr.update() return gr.update(value=image, label=f"Generated Image, seed = {seed}") # 이 함수는 사용되지 않으므로 제거해도 됩니다 # def generate_examples(id_image, control_image, prompt_text, seed, enable_realism, enable_anti_blur, model_version): # return generate_image(id_image, control_image, prompt_text, seed, 864, 1152, 3.5, 30, 1.0, 0.0, 1.0, enable_realism, enable_anti_blur, model_version) with gr.Blocks() as demo: session_state = gr.State({}) default_model_version = "v1.0" with gr.Row(): with gr.Column(scale=3): with gr.Row(): ui_id_image = gr.Image(label="Identity Image", type="pil", scale=3, height=370, min_width=100) with gr.Column(scale=2, min_width=100): ui_control_image = gr.Image(label="Control Image [Optional]", type="pil", height=370, min_width=100) ui_prompt_text = gr.Textbox(label="Prompt", value="Portrait, 4K, high quality, cinematic") ui_model_version = gr.Dropdown( label="Model Version", choices=[ModelVersion.STAGE_1, ModelVersion.STAGE_2], value=ModelVersion.DEFAULT_VERSION, ) ui_btn_generate = gr.Button("Generate") with gr.Accordion("Advanced", open=False): with gr.Row(): ui_num_steps = gr.Number(label="num steps", value=30) ui_seed = gr.Number(label="seed (0 for random)", value=0) with gr.Row(): ui_width = gr.Number(label="width", value=864) ui_height = gr.Number(label="height", value=1152) ui_guidance_scale = gr.Number(label="guidance scale", value=3.5, step=0.5) ui_infusenet_conditioning_scale = gr.Slider(minimum=0.0, maximum=1.0, value=1.0, step=0.05, label="infusenet conditioning scale") with gr.Row(): ui_infusenet_guidance_start = gr.Slider(minimum=0.0, maximum=1.0, value=0.0, step=0.05, label="infusenet guidance start") ui_infusenet_guidance_end = gr.Slider(minimum=0.0, maximum=1.0, value=1.0, step=0.05, label="infusenet guidance end") with gr.Accordion("LoRAs [Optional]", open=True): with gr.Row(): ui_enable_realism = gr.Checkbox(label="Enable realism LoRA", value=ENABLE_REALISM_DEFAULT) ui_enable_anti_blur = gr.Checkbox(label="Enable anti-blur LoRA", value=ENABLE_ANTI_BLUR_DEFAULT) with gr.Column(scale=2): image_output = gr.Image(label="Generated Image", interactive=False, height=550, format='png') # gr.Examples 부분을 완전히 제거했습니다 ui_btn_generate.click( generate_image, inputs=[ ui_id_image, ui_control_image, ui_prompt_text, ui_seed, ui_width, ui_height, ui_guidance_scale, ui_num_steps, ui_infusenet_conditioning_scale, ui_infusenet_guidance_start, ui_infusenet_guidance_end, ui_enable_realism, ui_enable_anti_blur, ui_model_version ], outputs=[image_output], concurrency_id="gpu" ) huggingface_hub.login(os.getenv('PRIVATE_HF_TOKEN')) download_models() init_pipeline(model_version=ModelVersion.DEFAULT_VERSION, enable_realism=ENABLE_REALISM_DEFAULT, enable_anti_blur=ENABLE_ANTI_BLUR_DEFAULT) # demo.queue() demo.launch() # demo.launch(server_name='0.0.0.0') # IPv4 # demo.launch(server_name='[::]') # IPv6