File size: 3,066 Bytes
8b0c8fb
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
import gradio as gr
from diffusers import StableDiffusionPipeline
import torch

# Initialize the Stable Diffusion pipeline
model_id = "sd-legacy/stable-diffusion-v1-5"

# Check if CUDA is available
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")

# Initialize pipeline with appropriate settings based on device
if device == "cuda":
    pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
else:
    pipe = StableDiffusionPipeline.from_pretrained(model_id)
    print("Running on CPU - using float32 precision")

pipe = pipe.to(device)

def generate_image(prompt, negative_prompt="", num_inference_steps=50, seed=None, guidance_scale=7.5):
    """
    Generate an image from a text prompt using Stable Diffusion
    """
    try:
        # Set the seed if provided
        if seed is not None and seed.strip():
            generator = torch.Generator(device).manual_seed(int(seed))
        else:
            generator = None

        # Generate the image
        image = pipe(
            prompt,
            negative_prompt=negative_prompt,
            num_inference_steps=num_inference_steps,
            guidance_scale=guidance_scale,
            generator=generator
        ).images[0]
        
        return image
    except Exception as e:
        print(f"Error generating image: {str(e)}")
        return None

# Create the Gradio interface
with gr.Blocks(title="Stable Diffusion Text-to-Image") as demo:
    gr.Markdown("# Stable Diffusion Text-to-Image Generator")
    gr.Markdown(f"Running on: {device.upper()}")
    gr.Markdown("Enter a text prompt to generate an image using Stable Diffusion")
    
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(
                label="Enter your prompt",
                placeholder="A beautiful sunset over mountains, digital art",
                lines=3
            )
            neg_prompt = gr.Textbox(
                label="Enter your negative prompt",
                value="ugly, tiling, poorly drawn hands, poorly drawn feet, poorly drawn face, out of frame,extra limbs, disfigured, deformed, body out of frame, bad anatomy, watermark, signature,cut off, low contrast, underexposed, overexposed, bad art, beginner, amateur, distorted face",
                lines=3,
                visible=False
            )
            steps = gr.Slider(
                minimum=2,
                maximum=50,
                value=20,
                step=1,
                label="Number of inference steps"
            )
            seed = gr.Textbox(
                label="Seed (optional)",
                placeholder="Leave empty for random seed"
            )
            generate_btn = gr.Button("Generate Image", variant='primary')
        
        with gr.Column():
            output = gr.Image(label="Generated Image")
    
    generate_btn.click(
        fn=generate_image,
        inputs=[prompt, neg_prompt, steps, seed],
        outputs=output
    )

if __name__ == "__main__":
    demo.launch(share=True)