image_generator / app.py
Eric2mangel's picture
Update app.py
ae91453 verified
Raw
History Blame Contribute Delete
1.08 kB
import gradio as gr
from diffusers import StableDiffusionPipeline
import torch
def generate_image(prompt):
try:
# Charger le modèle en précision float32 pour CPU
model_id = "runwayml/stable-diffusion-v1-5"
pipe = StableDiffusionPipeline.from_pretrained(
model_id,
torch_dtype=torch.float32, # Float32 pour compatibilité CPU
low_cpu_mem_usage=True, # Réduit l'usage mémoire
)
pipe = pipe.to("cpu") # Forcer explicitement CPU
# Générer l'image (réduire la taille pour économiser mémoire)
image = pipe(prompt, num_inference_steps=20, height=256, width=256).images[0]
return image
except Exception as e:
return f"Erreur : {str(e)}"
# Interface Gradio
with gr.Blocks() as demo:
gr.Markdown("# Image Generator with Stable Diffusion (CPU)")
with gr.Row():
prompt = gr.Textbox(label="Enter your prompt")
btn = gr.Button("Generate")
output = gr.Image()
btn.click(generate_image, inputs=prompt, outputs=output)
demo.launch()