File size: 1,564 Bytes
01518e3
 
 
3faffcb
ad059fb
 
01518e3
 
3faffcb
6e0d2cb
 
01518e3
3faffcb
6e0d2cb
 
3faffcb
6e0d2cb
01518e3
6e0d2cb
 
 
3faffcb
6e0d2cb
 
01518e3
6e0d2cb
 
 
 
01518e3
 
6e0d2cb
 
3faffcb
01518e3
 
 
3faffcb
 
 
 
 
 
01518e3
 
 
 
 
 
 
 
6e0d2cb
01518e3
6e0d2cb
01518e3
 
 
 
 
5efda0a
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
import os
import gradio as gr
import torch
from diffusers import DDPMPipeline
import warnings
warnings.filterwarnings('ignore')

# Constants
MODEL_ID = "google/ddpm-celebahq-256"
DEVICE = "cpu"  # Force CPU for better compatibility
DTYPE = torch.float32

def generate_image(steps=30):
    try:
        # Initialize pipeline with basic settings
        pipe = DDPMPipeline.from_pretrained(MODEL_ID)
        pipe = pipe.to(DEVICE)
        
        # Generate image
        with torch.inference_mode():
            image = pipe(
                batch_size=1,
                num_inference_steps=steps,
            ).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="Simple Image Generator") as demo:
    gr.Markdown("# 🎨 Simple Image Generator")
    gr.Markdown("Generate celebrity-like faces using DDPM")
    
    with gr.Row():
        with gr.Column():
            steps = gr.Slider(
                minimum=10,
                maximum=50,
                value=30,
                step=1,
                label="Steps"
            )
            
            generate_btn = gr.Button("🎨 Generate", variant="primary")
        
        with gr.Column():
            output_image = gr.Image(label="Generated Image", type="pil")
            
    generate_btn.click(
        fn=generate_image,
        inputs=[
            steps,
        ],
        outputs=output_image
    )

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