Spaces:
Runtime error
Runtime error
import torch | |
from diffusers import StableDiffusionPipeline | |
import gradio as gr | |
import os | |
import spaces | |
HF_TOKEN = os.getenv['HF_TOKEN'] | |
if torch.cuda.is_available(): | |
device = "cuda" | |
print("Using GPU") | |
else: | |
device = "cpu" | |
print("Using CPU") | |
# Initialize the pipeline and download the model | |
pipe = StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-3-medium", torch_dtype=torch.float16, use_auth_token=HF_TOKEN) | |
pipe.to(device) | |
# Define the image generation function | |
def generate_image(prompt): | |
image = pipe(prompt=prompt, num_inference_steps=1, guidance_scale=0.0).images[0] | |
return image | |
# Create the Gradio interface | |
interface = gr.Interface( | |
fn=generate_image, | |
inputs=gr.Textbox(placeholder="Insert Prompt Here..."), | |
outputs="image", | |
title="Stable Diffusion 3 Medium", | |
description="Made by [Nick088](https://linktr.ee/Nick088" | |
) | |
# Launch the interface | |
interface.launch(share=False) |