Nick088's picture
Update app.py
925b9c1 verified
raw
history blame
977 Bytes
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
@spaces.GPU(duration=60)
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)