| import torch |
| import gradio as gr |
| from diffusers import DiffusionPipeline |
|
|
|
|
| MODEL_ID = "stabilityai/stable-diffusion-xl-base-1.0" |
|
|
|
|
| |
| |
| |
|
|
| assert torch.cuda.is_available(), "CUDA GPU is required" |
|
|
| device = "cuda" |
|
|
| print("GPU:", torch.cuda.get_device_name(0)) |
| print( |
| "VRAM:", |
| round(torch.cuda.get_device_properties(0).total_memory / 1024**3, 2), |
| "GB" |
| ) |
|
|
|
|
| |
| torch.backends.cuda.matmul.allow_tf32 = True |
| torch.backends.cudnn.allow_tf32 = True |
|
|
|
|
| |
| |
| |
|
|
| pipe = DiffusionPipeline.from_pretrained( |
| MODEL_ID, |
| dtype=torch.float16, |
| variant="fp16", |
| use_safetensors=True, |
| ) |
|
|
| pipe = pipe.to(device) |
|
|
|
|
| |
| |
| |
|
|
| |
| pipe.unet.to(memory_format=torch.channels_last) |
| pipe.vae.to(memory_format=torch.channels_last) |
|
|
|
|
| |
| print("Compiling UNet...") |
|
|
| pipe.unet = torch.compile( |
| pipe.unet, |
| mode="reduce-overhead", |
| fullgraph=True |
| ) |
|
|
| print("UNet compilation enabled.") |
|
|
|
|
| |
| |
| |
|
|
| def generate_image(prompt): |
|
|
| if not prompt or not prompt.strip(): |
| raise gr.Error("Please enter a prompt.") |
|
|
| print("Generating image...") |
| print("Prompt:", prompt) |
|
|
| with torch.inference_mode(): |
|
|
| image = pipe( |
| prompt=prompt, |
| num_inference_steps=20, |
| guidance_scale=7.0, |
| height=768, |
| width=768, |
| ).images[0] |
|
|
| return image |
|
|
|
|
| |
| |
| |
|
|
| demo = gr.Interface( |
| fn=generate_image, |
|
|
| inputs=gr.Textbox( |
| label="Prompt", |
| placeholder="Describe the image you want...", |
| value=( |
| "a professional portrait of a confident woman " |
| "in a modern office, natural lighting, high quality" |
| ), |
| lines=3, |
| ), |
|
|
| outputs=gr.Image( |
| label="Generated Image", |
| type="pil", |
| ), |
|
|
| title="Stable Diffusion XL Image Generator", |
|
|
| description=( |
| "Fast SDXL generation using NVIDIA T4 GPU" |
| ), |
|
|
| examples=[ |
| [ |
| "a professional portrait of a confident woman " |
| "in a modern office, natural lighting, high quality" |
| ], |
| [ |
| "a futuristic city at night, neon lights, " |
| "cinematic photography" |
| ], |
| [ |
| "a beautiful mountain landscape at sunrise, " |
| "highly detailed" |
| ], |
| [ |
| "a cute golden retriever sitting in a park, " |
| "realistic photography" |
| ], |
| ], |
| ) |
|
|
|
|
| if __name__ == "__main__": |
| demo.launch() |
