Spaces:
Paused
Paused
File size: 4,633 Bytes
1bafe30 9231de3 1bafe30 9231de3 1bafe30 9231de3 1bafe30 9231de3 1bafe30 9231de3 1bafe30 9231de3 1bafe30 9231de3 1bafe30 9231de3 1bafe30 |
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 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 |
import gradio as gr
import numpy as np
import spaces
import torch
import random
import os
from PIL import Image
# Import the pipeline from diffusers
from diffusers import FluxKontextPipeline
# --- Constants and Model Loading ---
MAX_SEED = np.iinfo(np.int32).max
# --- FIX 1: Handle Hugging Face Authentication ---
# This is a gated model. You must have access on Hugging Face and provide a token.
# 1. Visit https://huggingface.co/black-forest-labs/FLUX.1-Kontext-dev and accept the terms.
# 2. Get an access token from https://huggingface.co/settings/tokens
# 3. Add the token below or set it as an environment variable `HF_TOKEN`.
HF_TOKEN = os.getenv("HF_TOKEN", "YOUR_HUGGING_FACE_TOKEN") # Replace with your token
# Load the pretrained model
try:
if HF_TOKEN == "YOUR_HUGGING_FACE_TOKEN":
pipe = None
print("Warning: Hugging Face token not provided. Please replace 'YOUR_HUGGING_FACE_TOKEN' with your actual token.")
else:
pipe = FluxKontextPipeline.from_pretrained(
"black-forest-labs/FLUX.1-Kontext-dev",
torch_dtype=torch.bfloat16,
token=HF_TOKEN, # Use the token for authentication
).to("cuda")
except Exception as e:
pipe = None
print(f"Warning: Could not load the model on CUDA. GPU is required. Error: {e}")
# --- Core Inference Function for ChatInterface ---
@spaces.GPU
def chat_fn(message, chat_history, seed, randomize_seed, guidance_scale, steps, progress=gr.Progress(track_tqdm=True)):
"""
Performs image generation or editing based on user input from the chat interface.
"""
if pipe is None:
raise gr.Error("Model could not be loaded. This could be due to a missing Hugging Face token, no access to the model, or no CUDA-enabled GPU.")
prompt = message["text"]
files = message["files"]
if not prompt and not files:
raise gr.Error("Please provide a prompt and/or upload an image.")
if randomize_seed:
seed = random.randint(0, MAX_SEED)
generator = torch.Generator(device="cuda").manual_seed(int(seed))
input_image = None
if files:
print(f"Received image: {files[0]}")
input_image = Image.open(files[0]).convert("RGB")
image = pipe(
image=input_image,
prompt=prompt,
guidance_scale=guidance_scale,
num_inference_steps=steps,
generator=generator,
).images[0]
else:
print(f"Received prompt for text-to-image: {prompt}")
image = pipe(
prompt=prompt,
guidance_scale=guidance_scale,
num_inference_steps=steps,
generator=generator,
).images[0]
return image
# --- UI Definition using gr.ChatInterface ---
seed_slider = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42)
randomize_checkbox = gr.Checkbox(label="Randomize seed", value=False)
guidance_slider = gr.Slider(label="Guidance Scale", minimum=1.0, maximum=10.0, step=0.1, value=2.5)
steps_slider = gr.Slider(label="Steps", minimum=1, maximum=30, value=28, step=1)
# --- FIX 2: Correctly format the examples as a list of lists ---
# Format: [ [message_dict, seed, randomize, guidance, steps], ... ]
examples = [
[
{"text": "A cute robot reading a book", "files": []},
42, False, 2.5, 28
],
[
{"text": "change his shirt to a hawaiian shirt", "files": ["https://gradio-builds.s3.amazonaws.com/demo-files/chewbacca.png"]},
12345, False, 3.0, 25
],
[
{"text": "make it a wooden house, add a chimney", "files": ["https://gradio-builds.s3.amazonaws.com/demo-files/house.png"]},
54321, False, 2.0, 30
],
]
demo = gr.ChatInterface(
fn=chat_fn,
title="FLUX.1 Kontext [dev]",
description="""<p style='text-align: center;'>
A simple chat UI for the <b>FLUX.1 Kontext</b> model.
<br>
To edit an image, upload it and type your instructions (e.g., "Add a hat").
<br>
To generate an image, just type a prompt (e.g., "A photo of an astronaut on a horse").
<br>
Find the model on <a href='https://huggingface.co/black-forest-labs/FLUX.1-Kontext-dev' target='_blank'>Hugging Face</a>.
</p>""",
textbox=gr.MultimodalTextbox(
file_types=["image"],
placeholder="Type a prompt and/or upload an image...",
render=False
),
additional_inputs=[
seed_slider,
randomize_checkbox,
guidance_slider,
steps_slider
],
examples=examples, # Use the correctly formatted list
theme="soft"
)
if __name__ == "__main__":
demo.launch() |