File size: 3,643 Bytes
5933dfb
 
 
 
 
dcf1a60
5933dfb
dcf1a60
 
5933dfb
dcf1a60
 
 
5933dfb
dcf1a60
5933dfb
dcf1a60
5933dfb
dcf1a60
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
5933dfb
dcf1a60
 
 
 
 
 
 
 
 
5933dfb
dcf1a60
 
 
 
 
 
 
 
 
 
 
 
 
 
5933dfb
dcf1a60
 
5933dfb
 
dcf1a60
5933dfb
 
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
89
90
91
92
93
94
95
96
97
import gradio as gr
import random
import torch
from diffusers import DiffusionPipeline

# Device setup
device = "cuda" if torch.cuda.is_available() else "cpu"
model_repo_id_turbo = "stabilityai/sdxl-turbo"  # Stability AI Model
pipe_turbo = DiffusionPipeline.from_pretrained(model_repo_id_turbo, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32).to(device)

# Placeholder for ZB-Tech model
def load_zb_model():
    return gr.Interface.load("models/ZB-Tech/Text-to-Image")

# Inference function
def custom_infer(
    model_choice, prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps
):
    # Load the selected model
    if model_choice == "Faster image generation (suitable for CPUs)":
        model = load_zb_model()
        return model(prompt)
    else:
        default_negative_prompt = "no watermark, hezzy, blurry"
        combined_negative_prompt = f"{default_negative_prompt}, {negative_prompt}" if negative_prompt else default_negative_prompt
        
        if randomize_seed:
            seed = random.randint(0, np.iinfo(np.int32).max)
        
        generator = torch.Generator().manual_seed(seed)
        image = pipe_turbo(
            prompt=prompt,
            negative_prompt=combined_negative_prompt,
            guidance_scale=guidance_scale,
            num_inference_steps=num_inference_steps,
            width=width,
            height=height,
            generator=generator,
        ).images[0]
        return image, seed

# CSS for centering UI
css = """
#col-container {
    display: flex;
    flex-direction: column;
    align-items: center;
    justify-content: center;
    text-align: center;
    margin: 0 auto;
}
"""

# Gradio app
with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
        # App name and description
        gr.Markdown(
            """
            # AI-Powered Text-to-Image Generator  
            *Generate stunning images from text prompts using advanced AI models.*
            """
        )

        # Dropdown for model selection
        model_choice = gr.Dropdown(
            label="Select Model",
            choices=[
                "Faster image generation (suitable for CPUs)",
                "More customizable option (slower, suitable for GPUs)"
            ],
            value="Faster image generation (suitable for CPUs)",
        )

        # Input section
        prompt = gr.Textbox(label="Prompt", placeholder="Enter your prompt here...")
        with gr.Accordion("Advanced Settings", open=False):
            negative_prompt = gr.Textbox(label="Negative Prompt", placeholder="Enter a negative prompt here...")
            seed = gr.Slider(label="Seed", minimum=0, maximum=2147483647, step=1, value=0)
            randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
            width = gr.Slider(label="Width", minimum=256, maximum=1024, step=32, value=512)
            height = gr.Slider(label="Height", minimum=256, maximum=1024, step=32, value=512)
            guidance_scale = gr.Slider(label="Guidance Scale", minimum=0.0, maximum=10.0, step=0.1, value=7.5)
            num_inference_steps = gr.Slider(label="Inference Steps", minimum=1, maximum=50, step=1, value=25)

        # Output section
        result = gr.Image(label="Generated Image", type="pil")
        gr.Button("Generate").click(
            custom_infer,
            inputs=[model_choice, prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
            outputs=result
        )

# Launch app
if __name__ == "__main__":
    demo.launch()