File size: 2,814 Bytes
4f91ffe
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
683afc3
88d1237
3aadc38
0737dc8
68e88ea
3aadc38
4fbc46c
c1497a6
3aadc38
68e88ea
d8f1f69
3aadc38
 
 
 
d8f1f69
 
 
 
 
 
 
 
 
 
 
 
 
 
 
68e88ea
3aadc38
 
 
 
 
 
 
68e88ea
3aadc38
 
 
 
 
 
 
68e88ea
d8f1f69
3aadc38
68e88ea
3aadc38
d8f1f69
 
 
 
68e88ea
91a655a
 
68e88ea
d8f1f69
 
 
3aadc38
d8f1f69
3aadc38
 
68e88ea
 
 
 
 
 
 
 
d8f1f69
 
68e88ea
 
a4cc7b2
68e88ea
 
d8f1f69
68e88ea
 
 
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
98
99
100
101
102
103
104
105
import os
import requests

url = "https://huggingface.co/InstantX/SD3.5-Large-IP-Adapter/resolve/main/ip-adapter.bin"
file_path = "ip-adapter.bin"

# Check if the file already exists
if not os.path.exists(file_path):
    print("File not found, downloading...")
    response = requests.get(url, stream=True)
    with open(file_path, "wb") as file:
        for chunk in response.iter_content(chunk_size=1024):
            if chunk:
                file.write(chunk)
    print("Download completed!")
else:
    print("File already exists.")

from models.transformer_sd3 import SD3Transformer2DModel
import gradio as gr
import torch
from pipeline_stable_diffusion_3_ipa import StableDiffusion3Pipeline
import os
import spaces
from huggingface_hub import login
token = os.getenv("HF_TOKEN")
login(token=token)

# Model and Pipeline Setup

model_path = 'stabilityai/stable-diffusion-3.5-large'
ip_adapter_path = './ip-adapter.bin'
image_encoder_path = "google/siglip-so400m-patch14-384"

transformer = SD3Transformer2DModel.from_pretrained(
    model_path, subfolder="transformer", torch_dtype=torch.bfloat16
)

pipe = StableDiffusion3Pipeline.from_pretrained(
    model_path, transformer=transformer, torch_dtype=torch.bfloat16
).to("cuda")

pipe.init_ipadapter(
    ip_adapter_path=ip_adapter_path,
    image_encoder_path=image_encoder_path,
    nb_token=64,
)


# Load transformer and pipeline
transformer = SD3Transformer2DModel.from_pretrained(
    model_path, subfolder="transformer", torch_dtype=torch.bfloat16
)
pipe = StableDiffusion3Pipeline.from_pretrained(
    model_path, transformer=transformer, torch_dtype=torch.bfloat16
).to("cuda")

# Initialize IP Adapter
pipe.init_ipadapter(
    ip_adapter_path=ip_adapter_path,
    image_encoder_path=image_encoder_path,
    nb_token=64,
)


@spaces.GPU
def gui_generation(prompt, ref_img):
    """
    Generate images using Stable Diffusion 3.5
    """
    image = pipe(
        width=1024,
        height=1024,
        prompt=prompt,
        negative_prompt="lowres, low quality, worst quality",
        num_inference_steps=24,
        guidance_scale=5.0,
        generator=torch.Generator("cuda").manual_seed(42),
        clip_image=ref_img,
        ipadapter_scale=0.5,
    ).images[0]

    return image


# Create Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# Stable Diffusion 3.5 Image Generation")

    with gr.Row():
        prompt_box = gr.Textbox(label="Prompt", placeholder="Enter your image generation prompt")

    with gr.Row():
        ref_img = gr.Image(type="pil", label="Upload Reference Image")
    gallery = gr.Image(type="pil", label="Generated Image")

    generate_btn = gr.Button("Generate")

    generate_btn.click(
        fn=gui_generation,
        inputs=[prompt_box, ref_img],
        outputs=gallery
    )
demo.launch()