import json import numpy as np import math import csv import random import argparse import torch import os import torch.distributed as dist import gradio as gr from PIL import Image from torch.nn.parallel import DistributedDataParallel as DDP import spaces from accelerate.utils import set_seed from diffusion_pipeline.sd35_pipeline import StableDiffusion3Pipeline, FlowMatchEulerInverseScheduler from diffusion_pipeline.sdxl_pipeline import StableDiffusionXLPipeline from diffusers import BitsAndBytesConfig, SD3Transformer2DModel from diffusers import FlowMatchEulerDiscreteScheduler, DDIMInverseScheduler, DDIMScheduler from huggingface_hub import login import os login(token=os.getenv('HF_TOKEN')) device = torch.device('cuda') # Load models outside the function to avoid reloading every time def load_models(): # Load sd35 model nf4_config = BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16 ) model_nf4 = SD3Transformer2DModel.from_pretrained( "stabilityai/stable-diffusion-3.5-large", subfolder="transformer", quantization_config=nf4_config, torch_dtype=torch.bfloat16 ) pipe_sd35 = StableDiffusion3Pipeline.from_pretrained( "stabilityai/stable-diffusion-3.5-large", transformer=model_nf4, torch_dtype=torch.bfloat16, ) pipe_sd35.scheduler = FlowMatchEulerDiscreteScheduler.from_config(pipe_sd35.scheduler.config) inverse_scheduler_sd35 = FlowMatchEulerInverseScheduler.from_pretrained( "stabilityai/stable-diffusion-3.5-large", subfolder='scheduler' ) pipe_sd35.inv_scheduler = inverse_scheduler_sd35 # Load sdxl model pipe_sdxl = StableDiffusionXLPipeline.from_pretrained( "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, variant="fp16", use_safetensors=True ).to("cuda") pipe_sdxl.scheduler = DDIMScheduler.from_config(pipe_sdxl.scheduler.config) inverse_scheduler_sdxl = DDIMInverseScheduler.from_pretrained( "stabilityai/stable-diffusion-xl-base-1.0", subfolder='scheduler' ) pipe_sdxl.inv_scheduler = inverse_scheduler_sdxl return pipe_sd35, pipe_sdxl pipe_sd35, pipe_sdxl = load_models() @spaces.GPU(duration=360) def generate_image( model_name, seed, num_steps, guidance_scale, inv_cfg, w2s_guidance, end_timesteps, prompt, method, size, ): try: # 根据传入的参数生成图像 torch.cuda.empty_cache() dtype = torch.float16 set_seed(seed) # Select the appropriate pipeline if model_name == 'sd35': pipe = pipe_sd35 elif model_name == 'sdxl': pipe = pipe_sdxl else: raise ValueError("Invalid model name") pipe.to(device) pipe.enable_model_cpu_offload() os.system('huggingface-cli download sst12345/CoRe2 weights/sd35_noise_model.pth weights/sdxl_noise_model.pth --local-dir ./weights') # TODO: load noise model if method == 'core' or method == 'z-core': from diffusion_pipeline.refine_model import PromptSD35Net, PromptSDXLNet from diffusion_pipeline.lora import replace_linear_with_lora, lora_true if model_name == 'sd35': refine_model = PromptSD35Net() replace_linear_with_lora(refine_model, rank=64, alpha=1.0, number_of_lora=28) lora_true(refine_model, lora_idx=0) checkpoint = torch.load('./weights/weights/sd35_noise_model.pth', map_location='cpu') refine_model.load_state_dict(checkpoint) elif model_name == 'sdxl': refine_model = PromptSDXLNet() replace_linear_with_lora(refine_model, rank=48, alpha=1.0, number_of_lora=50) lora_true(refine_model, lora_idx=0) checkpoint = torch.load('./weights/weights/sdxl_noise_model.pth', map_location='cpu') refine_model.load_state_dict(checkpoint) refine_model = refine_model.to(torch.bfloat16) refine_model = refine_model.to(device) print("Load Lora Success") # 根据模型类型设置形状 if model_name == 'sdxl': shape = (1, 4, size // 8, size // 8) else: shape = (1, 16, size // 8, size // 8) start_latents = torch.randn(shape, dtype=dtype).to(device) # 根据方法选择生成图像 if model_name == 'sdxl': if method == 'core': output = pipe.core( prompt=prompt, guidance_scale=guidance_scale, num_inference_steps=num_steps, latents=start_latents, return_dict=False, refine_model=refine_model, lora_true=lora_true, end_timesteps=end_timesteps, w2s_guidance=w2s_guidance)[0][0] elif method == 'zigzag': output = pipe.zigzag( prompt=prompt, guidance_scale=guidance_scale, latents=start_latents, return_dict=False, num_inference_steps=num_steps, inv_cfg=inv_cfg)[0][0] elif method == 'z-core': output = pipe.z_core( prompt=prompt, guidance_scale=guidance_scale, num_inference_steps=num_steps, latents=start_latents, return_dict=False, refine_model=refine_model, lora_true=lora_true, end_timesteps=end_timesteps, w2s_guidance=w2s_guidance, inv_cfg=inv_cfg)[0][0] elif method == 'standard': output = pipe( prompt=prompt, guidance_scale=guidance_scale, latents=start_latents, return_dict=False, num_inference_steps=num_steps)[0][0] else: raise ValueError("Invalid method") else: if method == 'core': output = pipe.core( prompt=prompt, guidance_scale=guidance_scale, num_inference_steps=num_steps, latents=start_latents, max_sequence_length=512, return_dict=False, refine_model=refine_model, lora_true=lora_true, end_timesteps=end_timesteps, w2s_guidance=w2s_guidance)[0][0] elif method == 'zigzag': output = pipe.zigzag( prompt=prompt, max_sequence_length=512, guidance_scale=guidance_scale, latents=start_latents, return_dict=False, num_inference_steps=num_steps, inv_cfg=inv_cfg)[0][0] elif method == 'z-core': output = pipe.z_core( prompt=prompt, guidance_scale=guidance_scale, num_inference_steps=num_steps, latents=start_latents, return_dict=False, max_sequence_length=512, refine_model=refine_model, lora_true=lora_true, end_timesteps=end_timesteps, w2s_guidance=w2s_guidance)[0][0] elif method == 'standard': output = pipe( prompt=prompt, guidance_scale=guidance_scale, latents=start_latents, return_dict=False, max_sequence_length=512, num_inference_steps=num_steps)[0][0] else: raise ValueError("Invalid method") # 将生成的图像保存为临时文件并返回 output_path = f'{model_name}_{method}.png' output.save(output_path) return output_path except Exception as e: print(f"An error occurred: {e}") return None if __name__ == '__main__': # 创建Gradio接口 iface = gr.Interface( fn=generate_image, inputs=[ gr.Dropdown(choices=['sdxl', 'sd35'], value='sdxl', label="Model"), # 设置默认模型为 'sdxl' gr.Slider(minimum=1, maximum=1000000, value=1, label="seed"), # 设置默认种子为 1 gr.Slider(minimum=1, maximum=100, value=50, label="Inference Steps"), # 设置默认推理步数为 50 gr.Slider(minimum=1, maximum=10, value=5.5, label="CFG"), # 设置默认CFG为 5.5 gr.Slider(minimum=-10, maximum=10, value=-1, label="Inverse CFG"), # 设置默认逆CFG为 -1 gr.Slider(minimum=1, maximum=3.5, value=2.5, label="W2S Guidance"), # 设置默认W2S指导为 2.5 gr.Slider(minimum=1, maximum=100, value=50, label="End Timesteps"), # 设置默认结束时间步为 50 gr.Textbox(label="Prompt"), # 文本框没有默认值 gr.Dropdown(choices=['standard', 'core', 'zigzag', 'z-core'], value='core', label="Method"), # 设置默认方法为 'core' gr.Slider(minimum=1024, maximum=2048, value=1024, label="Size") # 设置默认大小为 1024 ], outputs=gr.Image(type="filepath"), # 修改了type参数 title="Image Generation with CoRe^2", ) iface.launch(share=True)