Spaces:
Paused
Paused
import gradio as gr | |
from diffusers import AutoPipelineForText2Image, AutoPipelineForImage2Image | |
from diffusers.utils import load_image | |
import torch | |
# Clear CUDA cache | |
torch.cuda.empty_cache() | |
# Set environment variable for memory fragmentation | |
import os | |
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:128' | |
device = "cuda" if torch.cuda.is_available() else "cpu" | |
pipes = { | |
"txt2img": AutoPipelineForText2Image.from_pretrained("stabilityai/sdxl-turbo", torch_dtype=torch.float16, variant="fp16").to(device), | |
"img2img": AutoPipelineForImage2Image.from_pretrained("stabilityai/sdxl-turbo", torch_dtype=torch.float16, variant="fp16").to(device) | |
} | |
if device == "cpu": | |
pipes["txt2img"].enable_model_cpu_offload() | |
pipes["img2img"].enable_model_cpu_offload() | |
def run(prompt, image): | |
try: | |
print(f"prompt={prompt}, image={image}") | |
if image is None: | |
return pipes["txt2img"](prompt=prompt, num_inference_steps=1, guidance_scale=0.0).images[0] | |
else: | |
image = image.resize((512,512)) | |
print(f"img2img image={image}") | |
return pipes["img2img"](prompt, image=image, num_inference_steps=2, strength=0.5, guidance_scale=0.0).images[0] | |
except RuntimeError as e: | |
if "CUDA out of memory" in str(e): | |
print("CUDA out of memory. Trying to clear cache.") | |
torch.cuda.empty_cache() | |
# Consider additional fallback strategies here | |
else: | |
raise e | |
demo = gr.Interface( | |
run, | |
inputs=[ | |
gr.Textbox(label="Prompt"), | |
gr.Image(type="pil") | |
], | |
outputs=gr.Image(width=512, height=512), | |
live=True | |
) | |
demo.launch() | |