Spaces:
Build error
Build error
| # -*- coding: UTF-8 -*- | |
| """ | |
| @Time : 30/05/2025 19:24 | |
| @Author : xiaoguangliang | |
| @File : unconditional_diffusion_inference.py | |
| @Project : Faice_text2face | |
| """ | |
| import torch | |
| import random | |
| import numpy as np | |
| from diffusers import DDPMPipeline | |
| from accelerate import Accelerator | |
| import gradio as gr | |
| import spaces | |
| import PIL.Image | |
| from loguru import logger | |
| from utils import timer | |
| model_path = 'Ngene787/Faice_unconditional_diffusion' | |
| if torch.backends.mps.is_available(): | |
| accelerator = Accelerator(gradient_accumulation_steps=1) | |
| else: | |
| accelerator = Accelerator(mixed_precision="fp16", gradient_accumulation_steps=1) | |
| logger.info("Loading model ...") | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| if torch.cuda.is_available(): | |
| torch_dtype = torch.float16 | |
| else: | |
| torch_dtype = torch.float32 | |
| pipe = DDPMPipeline.from_pretrained(model_path, torch_dtype=torch_dtype, | |
| low_cpu_mem_usage=True | |
| ) | |
| pipe = pipe.to(device) | |
| pipe = accelerator.prepare(pipe) | |
| # Enable memory-efficient attention | |
| # pipe.enable_xformers_memory_efficient_attention() | |
| MAX_SEED = np.iinfo(np.int32).max | |
| def inference_unconditional(seed, | |
| randomize_seed=False, | |
| num_inference_steps=20, | |
| progress=gr.Progress(track_tqdm=True), ): | |
| if randomize_seed: | |
| seed = random.randint(0, MAX_SEED) | |
| generator = torch.Generator().manual_seed(seed) | |
| logger.info('Generating image ...') | |
| with timer("inference"): | |
| image = pipe( | |
| batch_size=1, | |
| generator=generator, | |
| num_inference_steps=num_inference_steps, | |
| output_type="np", | |
| ).images[0] | |
| return image | |