app.py

import gradio as gr
from huggingface_hub import login
import os
import spaces
import torch
from diffusers import StableDiffusionXLPipeline
from PIL import Image

from ip_adapter import IPAdapterXL

token = os.getenv("HF_TOKEN")
login(token=token)


base_model_path = "stabilityai/stable-diffusion-xl-base-1.0"
image_encoder_path = "sdxl_models/image_encoder"
ip_ckpt = "sdxl_models/ip-adapter_sdxl.bin"
device = "cuda"

# load SDXL pipeline
pipe = StableDiffusionXLPipeline.from_pretrained(
    base_model_path,
    torch_dtype=torch.float16,
    add_watermarker=False,
)

# reduce memory consumption
pipe.enable_vae_tiling()
ip_model = IPAdapterXL(pipe, image_encoder_path, ip_ckpt, device, target_blocks=["up_blocks.0.attentions.1"])


@spaces.GPU
def generate_image(prompt, reference_image, controlnet_conditioning_scale):
    image = Image.open(reference_image)
    image.resize((512, 512))
    images = ip_model.generate(pil_image=image,
                               prompt=prompt,
                               negative_prompt="",
                               scale=controlnet_conditioning_scale,
                               guidance_scale=5,
                               num_samples=1,
                               num_inference_steps=30,
                               seed=42,
                               # neg_content_prompt="a rabbit",
                               # neg_content_scale=0.5,
                               )

    return images[0]

# Set up Gradio interface
interface = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.Textbox(label="Prompt"),
        gr.Image( type= "filepath",label="Reference Image (Style)"),
        gr.Slider(label="Control Net Conditioning Scale", minimum=0, maximum=1.0, step=0.1, value=0.6),
    ],
    outputs="image",
    title="Image Generation with Stable Diffusion 3 medium and ControlNet",
    description="Generates an image based on a text prompt and a reference image using Stable Diffusion 3 medium with ControlNet."

)

interface.launch()