File size: 4,761 Bytes
6026000 |
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 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 |
import gradio as gr
import spaces
import torch
from model import ModelHandler
from generator import Generator
from config import Config
# 1. Initialize Models Globally
print("Initializing Application...")
handler = ModelHandler()
handler.load_models()
gen = Generator(handler)
# 2. Define GPU-enabled Inference Function
@spaces.GPU(duration=20)
def process_img(
image,
prompt,
negative_prompt,
cfg_scale, # <-- RE-ENABLED
steps,
img_strength,
depth_strength,
edge_strength,
seed
):
if image is None:
raise gr.Error("Please upload an image first.")
try:
print("--- Starting Generation ---")
result = gen.predict(
image,
prompt,
negative_prompt=negative_prompt,
guidance_scale=cfg_scale, # <-- RE-ENABLED
num_inference_steps=steps,
img2img_strength=img_strength,
depth_strength=depth_strength,
lineart_strength=edge_strength,
seed=seed
)
print("--- Generation Complete ---")
return result
except Exception as e:
print(f"Error during generation: {e}")
raise gr.Error(f"An error occurred: {str(e)}")
# 3. Build Gradio Interface
with gr.Blocks(title="Face To Style", theme=gr.themes.Soft()) as demo:
gr.Markdown(
"""
# 🎮 Face to Style
Upload any image. If there is a face, we'll keep the identity. If not, we'll stylize the scene!
"""
)
with gr.Row():
with gr.Column(scale=2):
input_img = gr.Image(type="pil", label="Input Image")
prompt = gr.Textbox(
label="Prompt (Optional)",
placeholder="Leave empty for auto-captioning...",
info=f"The trigger words '{Config.STYLE_TRIGGER}' are added automatically."
)
negative_prompt = gr.Textbox(
label="Negative Prompt (Optional)",
placeholder="e.g., blurry, text, watermark, bad art...",
value=Config.DEFAULT_NEGATIVE_PROMPT
)
with gr.Accordion("Advanced Settings", open=False):
seed = gr.Number(
label="Seed",
value=-1,
info="-1 for random",
precision=0
)
# --- RE-ENABLED CFG/GUIDANCE SLIDER ---
cfg_scale = gr.Slider(
elem_id="cfg_scale",
minimum=1.0,
maximum=10.0, # Range for TCD+Style
step=0.1,
value=Config.CGF_SCALE, # Default 4.0
label="Style Strength (Guidance)"
)
steps = gr.Slider(
elem_id="steps",
minimum=1,
maximum=20,
step=1,
value=8, # TCD default
label="Steps Number"
)
img_strength = gr.Slider(
elem_id="img_strength",
minimum=0.1,
maximum=1.0,
step=0.05,
value=Config.IMG_STRENGTH,
label="Image Strength (Img2Img)"
)
depth_strength = gr.Slider(
elem_id="depth_strength",
minimum=0.0,
maximum=1.0,
step=0.05,
value=Config.DEPTH_STRENGTH,
label="DepthMap Strength"
)
edge_strength = gr.Slider(
elem_id="edge_strength",
minimum=0.0,
maximum=1.0,
step=0.05,
value=Config.EDGE_STRENGTH,
label="EdgeMap Strength (LineArt)"
)
run_btn = gr.Button("Generate", variant="primary")
with gr.Column(scale=1):
output_img = gr.Image(label="Styled Result")
# Event Handler
all_inputs = [
input_img,
prompt,
negative_prompt,
cfg_scale, # <-- RE-ENABLED
steps,
img_strength,
depth_strength,
edge_strength,
seed
]
run_btn.click(
fn=process_img,
inputs=all_inputs,
outputs=[output_img]
)
# 4. Launch the App
if __name__ == "__main__":
demo.queue(max_size=20, api_open=True)
demo.launch(
server_name="0.0.0.0",
server_port=7860,
show_api=True
) |