Spaces:
Sleeping
Sleeping
File size: 2,986 Bytes
88b7650 2ffa99a 88b7650 2ffa99a 88b7650 2ffa99a 88b7650 2ffa99a 88b7650 2ffa99a 88b7650 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 |
import gradio as gr
import torch
from diffusers import AutoPipelineForText2Image
# 2. Load the model
# We remove the .to("cuda") from the end of this line.
print("Loading model...")
pipe = AutoPipelineForText2Image.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0",
torch_dtype=torch.float16,
variant="fp16",
use_safetensors=True
)
print("Model partially loaded.")
# ### CHANGE HERE ###
# This is the magic line that enables memory-efficient offloading.
pipe.enable_model_cpu_offload()
print("Model loaded with CPU offloading.")
# 3. Define the image generation function (This part remains the same)
def generate_image(prompt, negative_prompt, steps, guidance):
print(f"Generating image for prompt: {prompt}")
# The 'pipe' object does all the work. We pass it the prompt and other parameters.
# num_inference_steps controls how many steps the model takes to generate the image.
# guidance_scale controls how much the model follows the prompt.
image = pipe(
prompt=prompt,
negative_prompt=negative_prompt,
num_inference_steps=steps,
guidance_scale=guidance
).images[0]
# The output is a PIL Image object, which Gradio can display directly.
return image
# 4. Create the Gradio interface (This part remains the same)
with gr.Blocks(theme=gr.themes.Default(primary_hue="blue", secondary_hue="sky")) as demo:
gr.Markdown("# 🖼️ Stable Diffusion XL Text-to-Image")
gr.Markdown("Enter a text prompt and see the magic of AI-powered image generation!")
with gr.Row():
with gr.Column(scale=4):
prompt_input = gr.Textbox(label="Your Prompt", placeholder="An astronaut riding a horse on Mars, cinematic, 4k")
negative_prompt_input = gr.Textbox(label="Negative Prompt", placeholder="low quality, blurry, watermark, text")
submit_btn = gr.Button("Generate Image", variant="primary")
with gr.Column(scale=1):
steps_slider = gr.Slider(minimum=10, maximum=100, value=25, step=1, label="Inference Steps")
guidance_slider = gr.Slider(minimum=0, maximum=20, value=7.5, step=0.1, label="Guidance Scale")
output_image = gr.Image(label="Generated Image", width=768, height=768)
gr.Examples(
examples=[
["A majestic lion wearing a crown, photorealistic", "cartoon, drawing", 30, 8],
["A beautiful cityscape at sunset, painted by Van Gogh", "blurry, modern", 25, 7.5],
["A cute robot serving coffee in a futuristic cafe, 3D render", "text, humans", 35, 9],
],
inputs=[prompt_input, negative_prompt_input, steps_slider, guidance_slider]
)
# 5. Connect the button to the function (This part remains the same)
submit_btn.click(
fn=generate_image,
inputs=[prompt_input, negative_prompt_input, steps_slider, guidance_slider],
outputs=output_image
)
# 6. Launch the application
demo.launch() |