File size: 1,806 Bytes
eb2afd5 | 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 | import torch
import gradio as gr
from diffusers import StableDiffusionPipeline
# Chemin vers ton modèle .safetensors
model_path = "/Users/arthurdufour/Documents/ComfyUI/models/checkpoints/v1-5-pruned-emaonly.safetensors"
# Charger le modèle directement (évite load_state_dict)
pipeline = StableDiffusionPipeline.from_single_file(model_path, torch_dtype=torch.float32)
# Vérification du backend MPS pour MacBook M3
device = "mps" if torch.backends.mps.is_available() else "cpu"
pipeline.to(device)
def generate_image(positive_prompt, negative_prompt, steps, seed):
torch.mps.empty_cache() # Nettoyage mémoire
generator = torch.manual_seed(int(seed))
try:
image = pipeline(
prompt=positive_prompt,
negative_prompt=negative_prompt if "negative_prompt" in pipeline.__call__.__code__.co_varnames else None,
num_inference_steps=int(steps),
width=512,
height=512,
generator=generator
).images[0]
except Exception as e:
return f"Erreur : {str(e)}"
return image
# Interface Gradio
with gr.Blocks() as demo:
gr.Markdown("## Génération d'images Stable Diffusion (MPS)")
with gr.Row():
prompt_input = gr.Textbox(label="Prompt Positif", value="a horse")
negative_input = gr.Textbox(label="Prompt Négatif", value="text, watermark")
with gr.Row():
steps_slider = gr.Slider(1, 50, 20, step=1, label="Nombre de Steps")
seed_input = gr.Number(value=580029479038533, label="Seed")
output_image = gr.Image(label="Image Générée")
generate_button = gr.Button("Générer")
generate_button.click(generate_image, inputs=[prompt_input, negative_input, steps_slider, seed_input], outputs=output_image)
# Lancer l'interface
demo.launch()
|