import gradio as gr import numpy as np import random import torch import spaces # 사용할 파이프라인들을 모두 import 합니다. from diffusers import DiffusionPipeline # 'pipelines' 폴더가 있다고 가정합니다. from pipelines.pipeline_tag_stablediffusion import StableDiffusionTangentialDecomposedPipeline from pipelines.pipeline_tag_stablediffusion3 import StableDiffusion3TangentialDecomposedPipeline from pipelines.pipeline_tag_stablediffusionXL import StableDiffusionXLTangentialDecomposedPipeline # --- 설정 --- MODEL_MAP = { "SD 1.5": "runwayml/stable-diffusion-v1-5", "SD 2.1": "stabilityai/stable-diffusion-2-1", "SDXL": "stabilityai/stable-diffusion-xl-base-1.0", "SD 3": "stabilityai/stable-diffusion-3-medium-diffusers", } RESOLUTION_MAP = { "SD 1.5": 512, "SD 2.1": 768, "SDXL": 1024, "SD 3": 1024 } SEED_MAP = { "SD 1.5": 850728, "SD 2.1": 944905, "SDXL": 450040818, "SD 3": 282386105 } TAG_SCALE_MAP = { "SD 1.5": 1.15, # 기본값 "SD 2.1": 1.15, # 기본값 "SDXL": 1.20, "SD 3": 1.08 } PIPELINE_MAP = { "SD 1.5": StableDiffusionTangentialDecomposedPipeline, "SD 2.1": StableDiffusionTangentialDecomposedPipeline, "SDXL": StableDiffusionXLTangentialDecomposedPipeline, "SD 3": StableDiffusion3TangentialDecomposedPipeline, } device = "cuda" if torch.cuda.is_available() else "cpu" torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32 pipe = None current_model_id = None MAX_SEED = np.iinfo(np.int32).max # --- 함수 --- def load_pipeline(model_name, progress): global pipe, current_model_id model_id = MODEL_MAP[model_name] pipeline_class = PIPELINE_MAP[model_name] progress(0, desc=f"Loading model: {model_id} with {pipeline_class.__name__}...") if model_name == "SD 3": pipe = pipeline_class.from_pretrained(model_id, torch_dtype=torch_dtype, text_encoder_3=None, tokenizer_3=None,) else: pipe = pipeline_class.from_pretrained(model_id, torch_dtype=torch_dtype) pipe = pipe.to(device) current_model_id = model_id progress(1) def update_model_defaults(model_name): """모델 선택에 따라 해상도, 시드, 랜덤 시드 체크박스, TAG Scale을 업데이트합니다.""" res = RESOLUTION_MAP[model_name] seed_val = SEED_MAP[model_name] tag_scale_val = TAG_SCALE_MAP[model_name] return ( gr.update(value=res), gr.update(value=res), gr.update(value=seed_val), gr.update(value=False), # 'Randomize seed' 체크 해제 gr.update(value=tag_scale_val), # TAG Scale 업데이트 ) @spaces.GPU def infer( model_name, seed, randomize_seed, width, height, guidance_scale, # 사용자 지정 TAG Scale num_inference_steps, guidance_start_timestep, guidance_end_timestep, progress=gr.Progress(track_tqdm=True), ): global pipe, current_model_id model_id = MODEL_MAP[model_name] if model_id != current_model_id: gr.Info(f"Changing model to {model_name}. Please wait...") load_pipeline(model_name, progress) if randomize_seed: seed = random.randint(0, MAX_SEED) generator_custom = torch.Generator(device=device).manual_seed(int(seed)) generator_fixed = torch.Generator(device=device).manual_seed(int(seed)) unconditional_prompt = "" # 첫 번째 이미지 (사용자 지정 TAG Scale) 생성 image_custom_scale = pipe( prompt=unconditional_prompt, guidance_scale=0., num_inference_steps=num_inference_steps, width=width, height=height, generator=generator_custom, sta_tpd=guidance_start_timestep, end_tpd=guidance_end_timestep, t_guidance_scale=guidance_scale ).images[0] fixed_tag_value = 1.0 image_fixed_scale = pipe( prompt=unconditional_prompt, guidance_scale=0., num_inference_steps=num_inference_steps, width=width, height=height, generator=generator_fixed, sta_tpd=guidance_start_timestep, end_tpd=guidance_end_timestep, t_guidance_scale=fixed_tag_value ).images[0] return [image_fixed_scale, image_custom_scale], seed # --- UI 구성 (Gradio) --- css = """ #col-container { margin: 0 auto; max-width: 720px; } """ with gr.Blocks(css=css) as demo: with gr.Column(elem_id="col-container"): gr.Markdown("# Tangential Amplifying Guidance Demo") model_selector = gr.Dropdown( label="Select Model", choices=list(MODEL_MAP.keys()), value="SDXL" ) with gr.Row(): prompt = gr.Text( label="Prompt (Disabled)", show_label=True, max_lines=1, placeholder="Unconditional generation mode. This input is ignored.", container=True, interactive=False, ) run_button = gr.Button("Run", scale=0, variant="primary") # --- 2. gr.ImageSlider 컴포넌트로 변경 --- result_slider = gr.ImageSlider( label="Result Comparison (Fixed Scale vs. Your Scale)", show_label=True ) with gr.Accordion("Advanced Settings", open=False): seed = gr.Slider( label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=SEED_MAP["SDXL"] ) randomize_seed = gr.Checkbox(label="Randomize seed", value=False) with gr.Row(): width = gr.Slider( label="Width", minimum=256, maximum=1024, step=64, value=RESOLUTION_MAP["SDXL"] ) height = gr.Slider( label="Height", minimum=256, maximum=1024, step=64, value=RESOLUTION_MAP["SDXL"] ) with gr.Row(): guidance_scale = gr.Slider( label="TAG Scale", minimum=1.0, maximum=1.3, step=0.01, value=TAG_SCALE_MAP["SDXL"], ) num_inference_steps = gr.Slider( label="Inference Steps", minimum=20, maximum=50, step=1, value=50 ) with gr.Row(): guidance_start_timestep = gr.Slider( label="Guidance Start Timestep", minimum=0, maximum=1000, step=1, value=999 ) guidance_end_timestep = gr.Slider( label="Guidance End Timestep", minimum=0, maximum=1000, step=1, value=0 ) # --- 이벤트 리스너 --- model_selector.change( fn=update_model_defaults, inputs=[model_selector], outputs=[width, height, seed, randomize_seed, guidance_scale], ) # --- 3. outputs를 ImageSlider 컴포넌트로 지정 --- run_button.click( fn=infer, inputs=[ model_selector, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, guidance_start_timestep, guidance_end_timestep, ], outputs=[result_slider, seed], ) if __name__ == "__main__": demo.launch(debug=True)