Hugging Face's logo

Spaces:
dmpetrov
/
shapewords
Runtime error

App Files Community
shapewords / app.py
dmpetrov's picture
dmpetrov
better check
1725499
raw
history blame contribute delete
50.9 kB
 """
ShapeWords: Guiding Text-to-Image Synthesis with 3D Shape-Aware Prompts
=======================================================================
A Gradio web interface for the ShapeWords paper, allowing users to generate
images guided by 3D shape information.
Author: Melinos Averkiou
Date: 24 March 2025
Version: 1.5
Paper: "ShapeWords: Guiding Text-to-Image Synthesis with 3D Shape-Aware Prompts"
arXiv: https://arxiv.org/abs/2412.02912
Project Page: https://lodurality.github.io/shapewords/
Citation:
@misc{petrov2024shapewords,
title={ShapeWords: Guiding Text-to-Image Synthesis with 3D Shape-Aware Prompts},
author={Dmitry Petrov and Pradyumn Goyal and Divyansh Shivashok and Yuanming Tao and Melinos Averkiou and Evangelos Kalogerakis},
year={2024},
eprint={2412.02912},
archivePrefix={arXiv},
primaryClass={cs.CV},
url={https://arxiv.org/abs/2412.02912},
}
License: MIT License
Usage:
python app.py [--share]
This demo allows users to:
1. Select a 3D object category from ShapeNetCore
2. Choose a specific 3D shape using a slider or the navigation buttons (including a random shape button)
3. Enter a text prompt or pick a random one
4. Generate images guided by the selected 3D shape and the text prompt
The code is structured as a class and is compatible with Hugging Face ZeroGPU deployment.
"""
import os
import sys
import numpy as np
import torch
import gradio as gr
from PIL import Image, ImageFont, ImageDraw
from diffusers.utils import load_image
from diffusers import DPMSolverMultistepScheduler, StableDiffusionPipeline
import gdown
import argparse
import random
import spaces # for Hugging Face ZeroGPU deployment
import re
import plotly.graph_objects as go
from numpy.lib.user_array import container
import shutil
# Only for Hugging Face hosting - Add the Hugging Face cache to persistent storage to avoid downloading safetensors every time the demo sleeps and wakes up
os.environ['HF_HOME'] = '/data/.huggingface'
class ShapeWordsDemo:
# Constants
NAME2CAT = {
"chair": "03001627", "table": "04379243", "jar": "03593526", "skateboard": "04225987",
"car": "02958343", "bottle": "02876657", "tower": "04460130", "bookshelf": "02871439",
"camera": "02942699", "airplane": "02691156", "laptop": "03642806", "basket": "02801938",
"sofa": "04256520", "knife": "03624134", "can": "02946921", "rifle": "04090263",
"train": "04468005", "pillow": "03938244", "lamp": "03636649", "trash bin": "02747177",
"mailbox": "03710193", "watercraft": "04530566", "motorbike": "03790512",
"dishwasher": "03207941", "bench": "02828884", "pistol": "03948459", "rocket": "04099429",
"loudspeaker": "03691459", "file cabinet": "03337140", "bag": "02773838",
"cabinet": "02933112", "bed": "02818832", "birdhouse": "02843684", "display": "03211117",
"piano": "03928116", "earphone": "03261776", "telephone": "04401088", "stove": "04330267",
"microphone": "03759954", "bus": "02924116", "mug": "03797390", "remote": "04074963",
"bathtub": "02808440", "bowl": "02880940", "keyboard": "03085013", "guitar": "03467517",
"washer": "04554684", "bicycle": "02834778", "faucet": "03325088", "printer": "04004475",
"cap": "02954340", "phone": "02992529", "clock": "03046257", "helmet": "03513137",
"microwave": "03761084", "plant": "03991062"
}
PREDEFINED_PROMPTS = [
'a low poly 3d rendering of a [CATEGORY]',
'an aquarelle drawing of a [CATEGORY]',
'a photo of a [CATEGORY] on a beach',
'a charcoal drawing of a [CATEGORY]',
'a Hieronymus Bosch painting of a [CATEGORY]',
'a [CATEGORY] under a tree',
'A Kazimir Malevich painting of a [CATEGORY]',
'a vector graphic of a [CATEGORY]',
'a Claude Monet painting of a [CATEGORY]',
'a Salvador Dali painting of a [CATEGORY]',
'an Art Deco poster of a [CATEGORY]'
]
def __init__(self):
# Initialize class attributes
self.pipeline = None
self.shape2clip_model = None
self.text_encoder = None
self.tokenizer = None
self.category_embeddings = {}
self.category_counts = {}
self.available_categories = []
self.shape_thumbnail_cache = {} # Cache for shape thumbnails
self.CAT2NAME = {v: k for k, v in self.NAME2CAT.items()}
self.category_point_clouds = {}
# Initialize all models and data
self.initialize_models()
def initialize_models(self):
# device = DEVICE
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Using device: {device} in initialize_models")
# Download Shape2CLIP code if it doesn't exist
if not os.path.exists("shapewords_paper_code/geometry_guidance_models.py"):
shutil.rmtree("shapewords_paper_code/", ignore_errors=True)
print("Loading models file")
os.system("git clone https://github.com/lodurality/shapewords_paper_code.git")
# Import Shape2CLIP model
sys.path.append("./shapewords_paper_code")
from shapewords_paper_code.geometry_guidance_models import Shape2CLIP
# Initialize the pipeline
self.pipeline = StableDiffusionPipeline.from_pretrained(
"stabilityai/stable-diffusion-2-1-base",
torch_dtype=torch.float16 if device.type == "cuda" else torch.float32
)
self.pipeline.scheduler = DPMSolverMultistepScheduler.from_config(
self.pipeline.scheduler.config,
algorithm_type="sde-dpmsolver++"
)
self.text_encoder = self.pipeline.text_encoder
self.tokenizer = self.pipeline.tokenizer
# Look for Shape2CLIP checkpoint in multiple locations
checkpoint_paths = [
"./projection_model-0920192.pth",
"/data/projection_model-0920192.pth" # if using Hugging Face persistent storage look in a /data/ directory
]
checkpoint_found = False
checkpoint_path = None
for path in checkpoint_paths:
if os.path.exists(path):
checkpoint_path = path
print(f"Found Shape2CLIP checkpoint at: {checkpoint_path}")
checkpoint_found = True
break
# Download Shape2CLIP checkpoint if not found
if not checkpoint_found:
checkpoint_path = "projection_model-0920192.pth"
print("Downloading Shape2CLIP model checkpoint...")
gdown.download("https://drive.google.com/uc?id=1nvEXnwMpNkRts6rxVqMZt8i9FZ40KjP7", checkpoint_path, quiet=False) # download in same directory as app.py
print("Download complete")
# Initialize Shape2CLIP model
self.shape2clip_model = Shape2CLIP(depth=6, drop_path_rate=0.1, pb_dim=384)
self.shape2clip_model.load_state_dict(torch.load(checkpoint_path, map_location=device))
self.shape2clip_model.eval()
# Scan for available embeddings
self.scan_available_embeddings()
def scan_available_embeddings(self):
self.available_categories = []
self.category_counts = {}
# Try to find PointBert embeddings for all 55 ShapeNetCore shape categories
for category, cat_id in self.NAME2CAT.items():
possible_filenames = [
f"{cat_id}_pb_embs.npz",
f"embeddings/{cat_id}_pb_embs.npz",
f"/data/shapenet_pointbert_tokens/{cat_id}_pb_embs.npz" # if using Hugging Face persistent storage look in a /data/shapenet_pointbert_tokens directory
]
found_file = None
for filename in possible_filenames:
if os.path.exists(filename):
found_file = filename
break
if found_file:
try:
pb_data = np.load(found_file)
if 'ids' in pb_data:
count = len(pb_data['ids'])
else:
# Try to infer the correct keys
keys = list(pb_data.keys())
if len(keys) >= 1:
count = len(pb_data[keys[0]])
else:
count = 0
if count > 0:
self.available_categories.append(category)
self.category_counts[category] = count
print(f"Found {count} embeddings for category '{category}'")
except Exception as e:
print(f"Error loading embeddings for {category}: {e}")
# Sort categories alphabetically
self.available_categories.sort()
print(f"Found {len(self.available_categories)} categories with embeddings")
print(f"Available categories: {', '.join(self.available_categories)}")
# No embeddings found for any category - DEMO CANNOT RUN - but still load the interface with a default placeholder category, an error will be displayed when trying to generate images
if not self.available_categories:
self.available_categories = ["chair"] # Fallback
self.category_counts["chair"] = 50 # Default value
def load_category_embeddings(self, category):
if category in self.category_embeddings:
return self.category_embeddings[category]
if category not in self.NAME2CAT:
return None, []
cat_id = self.NAME2CAT[category]
# Check for different possible embedding filenames and locations
possible_filenames = [
f"{cat_id}_pb_embs.npz",
f"embeddings/{cat_id}_pb_embs.npz",
f"/data/shapenet_pointbert_tokens/{cat_id}_pb_embs.npz" # if using Hugging Face persistent storage look in a /data/shapenet_pointbert_tokens directory
]
# Find the first existing file
pb_emb_filename = None
for filename in possible_filenames:
if os.path.exists(filename):
pb_emb_filename = filename
print(f"Found embeddings file: {pb_emb_filename}")
break
if pb_emb_filename is None:
print(f"No embeddings found for {category}")
return None, []
# Load embeddings
try:
print(f"Loading embeddings from {pb_emb_filename}...")
pb_data = np.load(pb_emb_filename)
# Check for different key names in the NPZ file
if 'ids' in pb_data and 'embs' in pb_data:
pb_dict = dict(zip(pb_data['ids'], pb_data['embs']))
else:
# Try to infer the correct keys
keys = list(pb_data.keys())
if len(keys) >= 2:
# Assume first key is for IDs and second is for embeddings
pb_dict = dict(zip(pb_data[keys[0]], pb_data[keys[1]]))
else:
print("Unexpected embedding file format")
return None, []
all_ids = sorted(list(pb_dict.keys()))
print(f"Loaded {len(all_ids)} shape embeddings for {category}")
# Cache the results
self.category_embeddings[category] = (pb_dict, all_ids)
return pb_dict, all_ids
except Exception as e:
print(f"Error loading embeddings: {e}")
print(f"Exception details: {str(e)}")
return None, []
def load_category_point_clouds(self, category):
"""Load all point clouds for a category from a single NPZ file"""
if category not in self.NAME2CAT:
return None
cat_id = self.NAME2CAT[category]
# Cache to avoid reloading
if category in self.category_point_clouds:
return self.category_point_clouds[category]
# Check for different possible point cloud filenames
possible_filenames = [
f"{cat_id}.npz",
f"point_clouds/{cat_id}_clouds.npz",
f"/point_clouds/{cat_id}_clouds.npz",
f"/data/point_clouds/{cat_id}_clouds.npz" # For Hugging Face persistent storage
]
# Find the first existing file
pc_filename = None
for filename in possible_filenames:
if os.path.exists(filename):
pc_filename = filename
print(f"Found point cloud file: {pc_filename}")
break
if pc_filename is None:
print(f"No point cloud file found for category {category}")
return None
# Load point clouds
try:
print(f"Loading point clouds from {pc_filename}...")
pc_data_map = np.load(pc_filename, allow_pickle=False)
pc_data = {'ids': pc_data_map['ids'], 'clouds': pc_data_map['clouds']}
# Cache the loaded data
self.category_point_clouds[category] = pc_data
return pc_data
except Exception as e:
print(f"Error loading point clouds: {e}")
return None
def get_shape_preview(self, category, shape_idx):
"""Get a 3D point cloud visualization for a specific shape"""
if shape_idx is None or shape_idx < 0:
return None
# Get shape ID
pb_dict, all_ids = self.load_category_embeddings(category)
if pb_dict is None or not all_ids or shape_idx >= len(all_ids):
return None
shape_id = all_ids[shape_idx]
# Load all point clouds for this category
pc_data = self.load_category_point_clouds(category)
if pc_data is None:
# Fallback to image if point clouds not available
return self.get_shape_image_preview(category, shape_idx, shape_id)
# Extract point cloud for this specific shape
try:
# Get the arrays from the npz file
ids = pc_data['ids']
clouds = pc_data['clouds']
matching_indices = np.where(ids == shape_id)[0]
# Check number of matches
if len(matching_indices) == 0:
# No matches found - log error and fall back to image
print(f"Error: Shape ID {shape_id} not found in point cloud data")
return self.get_shape_image_preview(category, shape_idx, shape_id)
elif len(matching_indices) > 1:
# Multiple matches found - unexpected data issue - we will get the first one
print(f"Warning: Multiple matches ({len(matching_indices)}) found for Shape ID {shape_id}. Using first match.")
# Get the corresponding point cloud
matching_idx = matching_indices[0]
points = clouds[matching_idx]
# Create 3D visualization
fig = self.get_shape_pointcloud_preview(points, title=f"Shape #{shape_idx}")
return fig
except Exception as e:
print(f"Error extracting point cloud for {shape_id}: {e}")
return self.get_shape_image_preview(category, shape_idx, shape_id)
def get_shape_image_preview(self, category, shape_idx, shape_id):
"""Fallback to image preview if point cloud not available"""
try:
preview_image = self.get_ulip_image(shape_id)
preview_image = preview_image.resize((300, 300))
# Convert PIL image to plotly figure
fig = go.Figure()
# Need to convert PIL image to a format plotly can use
import io
import base64
# Convert PIL image to base64
buf = io.BytesIO()
preview_image.save(buf, format='PNG')
img_str = base64.b64encode(buf.getvalue()).decode('utf-8')
# Add image to figure
fig.add_layout_image(
dict(
source=f"data:image/png;base64,{img_str}",
xref="paper", yref="paper",
x=0, y=1,
sizex=1, sizey=1,
sizing="contain",
layer="below"
)
)
fig.update_layout(
title=f"Shape 2D Preview - 3D not available",
xaxis=dict(showgrid=False, zeroline=False, visible=False, range=[0, 1]),
yaxis=dict(showgrid=False, zeroline=False, visible=False, range=[0, 1], scaleanchor="x", scaleratio=1),
margin=dict(l=0, r=0, b=0, t=0),
plot_bgcolor='rgba(0,0,0,0)' # Transparent background
)
return fig
except Exception as e:
print(f"Error loading preview for {shape_id}: {e}")
# Create empty figure with error message
fig = go.Figure()
fig.update_layout(
title=f"Error loading Shape #{shape_idx}",
annotations=[dict(
text="Preview not available",
showarrow=False,
xref="paper", yref="paper",
x=0.5, y=0.5,
ont=dict(size=16, color="#E53935"), # Red error text
align="center"
)],
margin=dict(l=0, r=0, b=0, t=0, pad=0),
plot_bgcolor='rgba(0,0,0,0)' # Transparent background
)
return fig
def get_shape_pointcloud_preview(self, points, title=None):
"""Create a clean 3D point cloud visualization with Y as up axis"""
# Sample points for better performance (fewer points = smoother interaction)
sampled_points = points[::1] # Take every Nth point
# Create 3D scatter plot with fixed color
fig = go.Figure(data=[go.Scatter3d(
x=sampled_points[:, 0],
y=sampled_points[:, 1], # Use Z as Y (up axis)
z=sampled_points[:, 2], # Use Y as Z
mode='markers',
marker=dict(
size=2.5,
color='#4285F4', # Fixed blue color
opacity=1
)
)])
fig.update_layout(
title=None,
scene=dict(
# Remove all axes elements
xaxis=dict(visible=False, showticklabels=False, showgrid=False, zeroline=False, showline=False,
showbackground=False),
yaxis=dict(visible=False, showticklabels=False, showgrid=False, zeroline=False, showline=False,
showbackground=False),
zaxis=dict(visible=False, showticklabels=False, showgrid=False, zeroline=False, showline=False,
showbackground=False),
aspectmode='data' # Maintain data aspect ratio
),
# Eliminate margins
margin=dict(l=0, r=0, b=0, t=0, pad=0),
autosize=True,
# Control modebar appearance through layout
modebar=dict(
bgcolor='white',
color='#333',
orientation='v', # Vertical orientation
activecolor='#009688'
),
paper_bgcolor='rgba(0,0,0,0)', # Transparent background
)
# Better camera angle
fig.update_layout(
scene_camera=dict(
eye=dict(x=-1.5, y=0.5, z=-1.5),
up=dict(x=0, y=1, z=0), # Y is up
center=dict(x=0, y=0, z=0)
)
)
return fig
def get_ulip_image(self, guidance_shape_id, angle='036'):
shape_id_ulip = guidance_shape_id.replace('_', '-')
ulip_template = 'https://storage.googleapis.com/sfr-ulip-code-release-research/shapenet-55/only_rgb_depth_images/{}_r_{}_depth0001.png'
ulip_path = ulip_template.format(shape_id_ulip, angle)
try:
ulip_image = load_image(ulip_path).resize((512, 512))
return ulip_image
except Exception as e:
print(f"Error loading image: {e}")
return Image.new('RGB', (512, 512), color='gray')
def on_slider_change(self, shape_idx, category):
"""Update the preview when the slider changes"""
max_idx = self.category_counts.get(category, 0) - 1
# Get shape preview
shape_preview = self.get_shape_preview(category, shape_idx)
# Update counter text
counter_text = f"Shape {shape_idx} of {max_idx}"
return shape_preview, counter_text, shape_idx
def prev_shape(self, current_idx):
"""Go to previous shape"""
new_idx = max(0, current_idx - 1)
return new_idx
def next_shape(self, current_idx, category):
"""Go to next shape"""
max_idx = self.category_counts.get(category, 0) - 1
new_idx = min(max_idx, current_idx + 1)
return new_idx
def jump_to_start(self):
"""Jump to the first shape"""
return 0
def jump_to_end(self, category):
"""Jump to the last shape"""
max_idx = self.category_counts.get(category, 0) - 1
return max_idx
def random_shape(self, category):
"""Select a random shape from the category"""
max_idx = self.category_counts.get(category, 0) - 1
if max_idx <= 0:
return 0
# Generate random index
random_idx = random.randint(0, max_idx)
return random_idx
def random_prompt(self):
"""Select a random prompt from the predefined list"""
return random.choice(self.PREDEFINED_PROMPTS)
def on_category_change(self, category):
"""Update the slider and preview when the category changes"""
# Reset to the first shape
current_idx = 0
max_idx = self.category_counts.get(category, 0) - 1
# Get preview image
preview_image = self.get_shape_preview(category, current_idx)
# Update counter text
counter_text = f"Shape {current_idx} of {max_idx}"
# Need to update the slider range
new_slider = gr.Slider(
minimum=0,
maximum=max_idx,
step=1,
value=current_idx,
label="Shape Index"
)
return new_slider, current_idx, preview_image, counter_text
def get_guidance(self, test_prompt, category_name, guidance_emb):
print(test_prompt, category_name)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Using device: {device} in get_guidance")
prompt_tokens = torch.LongTensor(self.tokenizer.encode(test_prompt, padding='max_length')).to(device)
with torch.no_grad():
out = self.text_encoder(prompt_tokens.unsqueeze(0), output_attentions=True)
prompt_emb = out.last_hidden_state.detach().clone()
if len(guidance_emb.shape) == 1:
guidance_emb = torch.FloatTensor(guidance_emb).unsqueeze(0).unsqueeze(0)
else:
guidance_emb = torch.FloatTensor(guidance_emb).unsqueeze(0)
guidance_emb = guidance_emb.to(device)
eos_inds = torch.where(prompt_tokens.unsqueeze(0) == 49407)[1]
obj_word = category_name
obj_word_token = self.tokenizer.encode(obj_word)[-2]
chair_inds = torch.where(prompt_tokens.unsqueeze(0) == obj_word_token)[1]
eos_strength = 0.8
obj_strength = 1.0
self.shape2clip_model.eval()
with torch.no_grad():
guided_prompt_emb_cond = self.shape2clip_model(prompt_emb.float(), guidance_emb[:,:,:].float()).half()
guided_prompt_emb = guided_prompt_emb_cond.clone()
guided_prompt_emb[:,:1] = 0
guided_prompt_emb[:,:chair_inds] = 0
guided_prompt_emb[:,chair_inds] *= obj_strength
guided_prompt_emb[:,eos_inds+1:] = 0
guided_prompt_emb[:,eos_inds] *= eos_strength
guided_prompt_emb[:,chair_inds+1:eos_inds:] = 0
fin_guidance = guided_prompt_emb
return fin_guidance, prompt_emb
@spaces.GPU(duration=120)
def generate_images(self, prompt, category, selected_shape_idx, guidance_strength, seed):
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Using device: {device} in generate_images")
# Move models to gpu
if device.type == "cuda":
self.pipeline = self.pipeline.to(device)
self.shape2clip_model = self.shape2clip_model.to(device)
# Clear status text immediately
status = ""
# Replace [CATEGORY] with the selected category (case-insensitive)
category_pattern = re.compile(r'\[CATEGORY\]', re.IGNORECASE)
if re.search(category_pattern, prompt):
# Use re.sub for replacement to maintain the same casing pattern that was used
final_prompt = re.sub(category_pattern, category, prompt)
else:
# Fallback if user didn't use placeholder
final_prompt = f"{prompt} {category}"
status += self.create_status_message(
f"Warning! For better results, use [CATEGORY] in your prompt where you want '{category}' to appear, otherwise it is appended at the end of the prompt.",
"info"
)
error = False
# Check if prompt contains any other categories
for other_category in self.available_categories:
if re.search(r'\b' + re.escape(other_category) + r'\b', prompt, re.IGNORECASE):
status += self.create_status_message(
f"Error! Your prompt contains '{other_category}'. Please remove it and use [CATEGORY] instead.",
"error"
)
error = True
if error:
return [], status
# Load category embeddings if not already loaded
pb_dict, all_ids = self.load_category_embeddings(category)
if pb_dict is None or not all_ids:
status += self.create_status_message(
f"Error! Unable to load embeddings for {category}",
"error"
)
return [], status
# Ensure shape index is valid
if selected_shape_idx is None or selected_shape_idx < 0:
selected_shape_idx = 0
max_idx = len(all_ids) - 1
selected_shape_idx = max(0, min(selected_shape_idx, max_idx))
guidance_shape_id = all_ids[selected_shape_idx]
# Set generator
generator = torch.Generator(device=device).manual_seed(seed)
results = []
try:
# Generate base image (without guidance)
with torch.no_grad():
base_images = self.pipeline(
prompt=final_prompt,
num_inference_steps=50,
num_images_per_prompt=1,
generator=generator,
guidance_scale=7.5
).images
results.append((base_images[0], "Unguided Result"))
except Exception as e:
print(f"Error generating base image: {e}")
status += self.create_status_message(
f"Error! Unable to generate base image: {str(e)}",
"error"
)
return results, status
try:
# Get shape guidance embedding
pb_emb = pb_dict[guidance_shape_id]
out_guidance, prompt_emb = self.get_guidance(final_prompt, category, pb_emb)
except Exception as e:
print(f"Error generating guidance: {e}")
status += self.create_status_message(
f"Error! Unable to generate guidance: {str(e)}",
"error"
)
return results, status
try:
# Generate guided image
generator = torch.Generator(device=device).manual_seed(seed)
with torch.no_grad():
guided_images = self.pipeline(
prompt_embeds=prompt_emb + guidance_strength * out_guidance,
num_inference_steps=50,
num_images_per_prompt=1,
generator=generator,
guidance_scale=7.5
).images
results.append((guided_images[0], f"Guided Result (λ = {guidance_strength})"))
# Success status
status += self.create_status_message(
f"Success! Generated image guided by Shape #{selected_shape_idx} from category '{category}'.",
"success"
)
torch.cuda.empty_cache()
except Exception as e:
print(f"Error generating guided image: {e}")
status += self.create_status_message(
f"Error! Unable to generate guided image: {str(e)}",
"error"
)
return results, status
return results, status
def create_status_message(self, content, type_="info"):
# Define styles for different message types
styles = {
"info": {
"bg": "rgba(33, 150, 243, 0.15)",
"border": "#2196F3",
"icon": "ℹ️",
"title": "NOTE: "
},
"error": {
"bg": "rgba(244, 67, 54, 0.15)",
"border": "#F44336",
"icon": "❌",
"title": "ERROR: "
},
"success": {
"bg": "rgba(76, 175, 80, 0.15)",
"border": "#4CAF50",
"icon": "✅",
"title": "SUCCESS: "
},
"waiting": {
"bg": "rgba(255, 193, 7, 1)",
"border": "#FFC107",
"icon": "⏳",
"title": "PROCESSING: "
}
}
style = styles.get(type_, styles["info"])
font_weight = "bold" if type_ == "waiting" else "normal"
animation_style = "animation: pulse 1.5s infinite;" if type_ == "waiting" else ""
return f"""
<div style='
padding: 12px;
background-color: {style["bg"]};
border-left: 5px solid {style["border"]};
margin-bottom: 12px;
border-radius: 4px;
display: flex;
align-items: flex-start;
gap: 8px;
box-shadow: 0 1px 3px rgba(0,0,0,0.12);
font-weight: {font_weight};
{animation_style}
'>
<style>
@keyframes pulse {{
0%, 100% {{ opacity: 1; }}
50% {{ opacity: 0.7; }}
}}
</style>
<div style='font-size: 18px; line-height: 1.2;'>{style["icon"]}</div>
<div>{content}</div>
</div>
"""
def on_demo_load(self):
"""Function to ensure initial image is loaded when demo starts"""
default_category = "chair" if "chair" in self.available_categories else self.available_categories[0]
initial_img = self.get_shape_preview(default_category, 0)
return initial_img
def create_ui(self):
# Ensure chair is in available categories, otherwise use the first available
default_category = "chair" if "chair" in self.available_categories else self.available_categories[0]
with gr.Blocks(title="ShapeWords: Guiding Text-to-Image Synthesis with 3D Shape-Aware Prompts",
theme=gr.themes.Soft(
primary_hue="orange",
secondary_hue="blue",
font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui", "sans-serif"],
font_mono=[gr.themes.GoogleFont("IBM Plex Mono"), "ui-monospace", "Consolas", "monospace"],
),
css="""
/* Base styles */
.container { max-width: 1400px; margin: 0 auto; }
/* Title headers */
.title { text-align: center; font-size: 26px; font-weight: 600; margin-bottom: 3px; }
.subtitle { text-align: center; font-size: 16px; margin-bottom: 3px; }
.authors { text-align: center; font-size: 15px; margin-bottom: 3px; }
.affiliations { text-align: center; font-size: 13px; margin-bottom: 3px; }
/* Instructions Accordion */
button.instructions-accordion > span,
.instructions-accordion button > span {
font-size: 17px !important;
font-weight: 600 !important;
}
.instructions-accordion + div p,
.instructions-accordion + div li,
.instructions-text p,
.instructions-text li {
font-size: 14px !important;
}
/* Section Headers */
.step-header,.settings-header {
font-size: 18px;
font-weight: 600;
margin-top: 5px;
margin-bottom: 5px;
}
.sub-header {
margin-top: 5px;
margin-bottom: 5px;
padding-left: 5px;
}
/* Buttons for project page, paper, code etc*/
.buttons-container { margin: 0 auto 10px; }
.buttons-row { display: flex; justify-content: center; gap: 10px; flex-wrap: nowrap; }
.nav-button {
display: inline-block;
padding: 6px 12px;
background-color: #363636;
color: white !important;
text-decoration: none;
border-radius: 20px;
font-weight: 500;
font-size: 14px;
transition: background-color 0.2s;
text-align: center;
white-space: nowrap;
}
.nav-button:hover { background-color: #505050; }
.nav-button.disabled {
opacity: 0.6;
cursor: not-allowed;
}
/* Prompt design section elements */
.category-dropdown .wrap { font-size: 16px; }
.prompt-input { flex-grow: 1; }
.prompt-button {
align-self: center; /* Vertical centering */
margin-left: auto; /* Horizontal centering */
margin-right: auto;
display: block; /* Makes margins work for centering */
}
/* Shape selection section elements */
.shape-navigation {
display: flex;
justify-content: center;
align-items: center;
margin: 10px auto;
gap: 15px;
max-width: 320px;
}
.shape-navigation button {
min-width: 40px;
max-width: 60px;
width: auto;
padding: 6px 10px;
}
.nav-icon-btn { font-size: 18px; }
/* Generate button */
.generate-button {
font-size: 18px !important;
padding: 12px !important;
margin: 15px 0 !important;
background: linear-gradient(135deg, #f97316, #fb923c) !important;
}
/* Results section elements */
.results-gallery { min-height: 100px; max-height: 500px; display: flex; align-items: center; justify-content: center; }
.results-gallery .grid-container { display: flex; align-items: center; }
/* About section elements */
.about-section { font-size: 16px; margin-top: 40px; padding: 20px; border-top: 1px solid rgba(128, 128, 128, 0.2); }
/* Responsive adjustments for mobile mode*/
@media (max-width: 768px) {
.shape-navigation {
max-width: 100%;
gap: 5px;
}
.shape-navigation button {
min-width: 36px;
padding: 6px 0;
font-size: 16px;
}
.buttons-row {
gap: 5px;
}
.nav-button {
padding: 5px 8px;
font-size: 13px;
}
.results-gallery {
max-height: 320px;
}
}
/* Dark mode overrides */
@media (prefers-color-scheme: dark) {
.nav-button {
background-color: #505050;
}
.nav-button:hover {
background-color: #666666;
}
}
""") as demo:
# Header with title and links
gr.Markdown("# ShapeWords: Guiding Text-to-Image Synthesis with 3D Shape-Aware Prompts",
elem_classes="title")
gr.Markdown("### CVPR 2025", elem_classes="subtitle")
gr.Markdown(
"Dmitry Petrov<sup>1</sup>, Pradyumn Goyal<sup>1</sup>, Divyansh Shivashok<sup>1</sup>, Yuanming Tao<sup>1</sup>, Melinos Averkiou<sup>2,3</sup>, Evangelos Kalogerakis<sup>1,2,4</sup>",
elem_classes="authors")
gr.Markdown(
"<sup>1</sup>UMass Amherst <sup>2</sup>CYENS CoE <sup>3</sup>University of Cyprus <sup>4</sup>TU Crete",
elem_classes="affiliations")
# Navigation buttons
with gr.Row():
with gr.Column(scale=3):
pass # Empty space for alignment
with gr.Column(scale=2, elem_classes="buttons-container"):
gr.HTML("""
<div class="buttons-row">
<a href="https://arxiv.org/abs/2412.02912" target="_blank" class="nav-button">
arXiv
</a>
<a href="https://lodurality.github.io/shapewords/" target="_blank" class="nav-button">
Project Page
</a>
<a href="#" target="_blank" class="nav-button disabled">
Code
</a>
<a href="#" target="_blank" class="nav-button disabled">
Data
</a>
</div>
""")
with gr.Column(scale=3):
pass # Empty space for alignment
# Add instructions
with gr.Accordion("📋 Instructions", open=True, elem_classes="instructions-accordion"):
gr.Markdown("""
1️⃣ Select an shape category from the dropdown menu -- overall 55 categories. We recommend trying chair (default), car, lamp and bottle categories.
2️⃣ Create a text prompt using **[CATEGORY]** as a placeholder or use **"Random prompt"** button to select from a small set of pre-defined prompts
3️⃣ Adjust **guidance strength** to control shape influence. Use the default 0.9 value for best balance between prompt and shape adherence. Value of 0.0 corresponds to unguided result that is based just on input prompt.
4️⃣ (optional) Choose **random seed**. For a fixed combination of input prompt and random seed, unguided image will always be the same.
5️⃣ Choose **guidance 3D shape** using the slider, navigation or random shape buttons. Shapes come from ShapeNet dataset (~55K shapes across all categories)
6️⃣ Click **Generate Images** button at the bottom to create images that follow both your text prompt and the selected 3D shape geometry
""", elem_classes="instructions-text")
# Hidden field to store selected shape index
selected_shape_idx = gr.Number(value=0, visible=False)
# Prompt Design (full width)
with gr.Group():
gr.Markdown("### 📝 Prompt Design", elem_classes="step-header")
with gr.Row():
category = gr.Dropdown(
label="1️⃣ Shape Category",
choices=self.available_categories,
value=default_category,
container=True,
elem_classes="category-dropdown",
scale=2
)
prompt = gr.Textbox(
label="2️⃣ Text Prompt - Use [CATEGORY] as a placeholder, e.g. 'a [CATEGORY] under a tree'",
placeholder="an aquarelle drawing of a [CATEGORY]",
value="an aquarelle drawing of a [CATEGORY]",
lines=1,
scale=5,
elem_classes="prompt-input"
)
random_prompt_btn = gr.Button("🎲 Random\nPrompt",
size="lg",
scale=1,
elem_classes="prompt-button")
# Generation Settings (full width)
with gr.Group():
gr.Markdown("### ⚙️ Generation Settings", elem_classes="settings-header")
with gr.Row():
with gr.Column():
guidance_strength = gr.Slider(
minimum=0.0, maximum=1.0, step=0.1, value=0.9,
label="3️⃣ Guidance Strength (λ) - Higher λ = stronger shape adherence"
)
with gr.Column():
seed = gr.Slider(
minimum=0, maximum=10000, step=1, value=42,
label="4️⃣ Random Seed - (optional) Change for different variations"
)
# Middle section - Shape Selection and Results side by side
with gr.Row(equal_height=True):
# Left column - Shape Selection
with gr.Column():
with gr.Group():
gr.Markdown("### 🔍 Shape Selection", elem_classes="step-header")
shape_slider = gr.Slider(
minimum=0,
maximum=self.category_counts.get(default_category, 0) - 1,
step=1,
value=0,
label="5️⃣ Shape Index - Choose a 3D shape to guide image generation",
interactive=True
)
shape_counter = gr.Markdown(f"Shape 0 of {self.category_counts.get(default_category, 0) - 1}", elem_classes="sub-header")
current_shape_plot = gr.Plot(show_label=False)
# Navigation buttons - Icons only for better mobile compatibility
with gr.Row(elem_classes="shape-navigation"):
jump_start_btn = gr.Button("⏮️", size="sm", elem_classes="nav-icon-btn")
prev_shape_btn = gr.Button("◀️", size="sm", elem_classes="nav-icon-btn")
random_btn = gr.Button("🎲", size="sm", variant="secondary", elem_classes="nav-icon-btn")
next_shape_btn = gr.Button("▶️", size="sm", elem_classes="nav-icon-btn")
jump_end_btn = gr.Button("⏭️", size="sm", elem_classes="nav-icon-btn")
# Right column - Results
with gr.Column():
with gr.Group():
gr.Markdown("### 🖼️ Generated Results Preview", elem_classes="step-header")
gallery = gr.Gallery(
label="Results",
show_label=False,
elem_id="results_gallery",
columns=2,
elem_classes="results-gallery"
)
# Generate button (full width)
with gr.Row():
run_button = gr.Button(" 6️⃣ ✨ Generate Images guided by Selected Shape", variant="primary", size="lg",
elem_classes="generate-button")
# Status message (full width)
with gr.Row():
status_text = gr.HTML("", elem_classes="status-message")
# About section at the bottom of the page
with gr.Group(elem_classes="about-section"):
gr.Markdown("""
### About ShapeWords
ShapeWords incorporates target 3D shape information with text prompts to guide image synthesis.
### How It Works
1. Select an shape category from the dropdown menu -- overall 55 categories. We recommend trying chair (default), car, lamp and bottle categories.
2. Create a text prompt using **[CATEGORY]** as a placeholder or use **"Random prompt"** button to select from a small set of pre-defined prompts
3. Adjust **guidance strength** to control shape influence. Use the default 0.9 value for best balance between prompt and shape adherence. Value of 0.0 corresponds to unguided result that is based just on input prompt.
4. (optional) Choose **random seed**. For a fixed combination of input prompt and random seed, unguided image will always be the same.
5. Choose **guidance 3D shape** using the slider, navigation or random shape buttons. Shapes come from ShapeNet dataset (~55K shapes across all categories)
6. Click **Generate Images** button at the bottom to create images that follow both your text prompt and the selected 3D shape geometry
### Citation
```
@misc{petrov2024shapewords,
title={ShapeWords: Guiding Text-to-Image Synthesis with 3D Shape-Aware Prompts},
author={Dmitry Petrov and Pradyumn Goyal and Divyansh Shivashok and Yuanming Tao and Melinos Averkiou and Evangelos Kalogerakis},
year={2024},
eprint={2412.02912},
archivePrefix={arXiv},
primaryClass={cs.CV},
url={https://arxiv.org/abs/2412.02912},
}
```
""")
# Make sure the initial image is loaded when the demo starts
demo.load(
fn=self.on_demo_load,
inputs=None,
outputs=[current_shape_plot]
)
# Connect slider to update preview
shape_slider.change(
fn=self.on_slider_change,
inputs=[shape_slider, category],
outputs=[current_shape_plot, shape_counter, selected_shape_idx]
)
# Previous shape button
prev_shape_btn.click(
fn=self.prev_shape,
inputs=[selected_shape_idx],
outputs=[shape_slider]
)
# Next shape button
next_shape_btn.click(
fn=self.next_shape,
inputs=[selected_shape_idx, category],
outputs=[shape_slider]
)
# Jump to start button
jump_start_btn.click(
fn=self.jump_to_start,
inputs=None,
outputs=[shape_slider]
)
# Jump to end button
jump_end_btn.click(
fn=self.jump_to_end,
inputs=[category],
outputs=[shape_slider]
)
# Random shape button
random_btn.click(
fn=self.random_shape,
inputs=[category],
outputs=[shape_slider]
)
# Connect the random prompt button
random_prompt_btn.click(
fn=self.random_prompt,
inputs=[],
outputs=[prompt]
)
# Update the UI when category changes
category.change(
fn=self.on_category_change,
inputs=[category],
outputs=[shape_slider, selected_shape_idx, current_shape_plot, shape_counter]
)
# Update status text when generating
run_button.click(
fn=lambda: self.create_status_message("Generating images...", "waiting"),
inputs=None,
outputs=[status_text]
)
# Generate images when button is clicked
run_button.click(
fn=self.generate_images,
inputs=[prompt, category, selected_shape_idx, guidance_strength, seed],
outputs=[gallery, status_text]
)
return demo
# Main function and entry point
def main():
parser = argparse.ArgumentParser(description="ShapeWords Gradio Demo")
parser.add_argument('--share', action='store_true', help='Create a public link')
args = parser.parse_args()
# Create the demo app and UI
app = ShapeWordsDemo()
demo = app.create_ui()
demo.launch(share=args.share)
if __name__ == "__main__":
main()