File size: 1,616 Bytes
c61999a
 
07d15c1
a6709d2
fa80602
07d15c1
a6709d2
07d15c1
 
a6709d2
c61999a
 
07d15c1
c61999a
07d15c1
 
 
 
 
 
 
 
 
 
 
c61999a
07d15c1
 
c61999a
 
 
 
 
 
 
 
07d15c1
c61999a
 
07d15c1
 
c61999a
07d15c1
 
 
 
c61999a
 
07d15c1
 
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
import gradio as gr
import torch
from diffusers import StableDiffusionPipeline
import os

# Retrieve the Hugging Face token stored in Hugging Face Spaces secrets
HUGGINGFACE_TOKEN = os.getenv("keyss")

if not HUGGINGFACE_TOKEN:
    raise ValueError("Hugging Face token not found! Make sure it's added in Hugging Face Secrets.")

def image_generation(prompt):
    # Check if GPU is available
    device = "cuda" if torch.cuda.is_available() else "cpu"
    
    # Load the Stable Diffusion 3 pipeline
    pipeline = StableDiffusionPipeline.from_pretrained(
        "stabilityai/stable-diffusion-3-medium-diffusers",
        torch_dtype=torch.float16 if device == "cuda" else torch.float32,
        use_auth_token=HUGGINGFACE_TOKEN,  # Use the Hugging Face token for authentication
        text_encoder_3=None,
        tokenizer_3=None
    )
    
    # Enable efficient model execution
    pipeline.enable_model_cpu_offload()
    
    # Generate an image based on the prompt
    image = pipeline(
        prompt=prompt,
        negative_prompt="blurred, ugly, watermark, low resolution, blurry",
        num_inference_steps=40,
        height=1024,
        width=1024,
        guidance_scale=9.0
    ).images[0]
    
    return image

# Define the Gradio interface
interface = gr.Interface(
    fn=image_generation,
    inputs=gr.Textbox(lines=2, placeholder="Enter your Prompt..."),
    outputs=gr.Image(type="pil"),
    title="Image Creation using Stable Diffusion 3 Model",
    description="This application generates awesome images using the Stable Diffusion 3 model."
)

# Launch the Gradio app
interface.launch()