Spaces:
Paused
Paused
File size: 3,029 Bytes
fdd3761 8124057 c1ad0ba 37abfc2 b2ef34b 4e7c237 38d816c 37abfc2 b2ef34b 37abfc2 d3a3bf1 37abfc2 d3a3bf1 37abfc2 d3a3bf1 37abfc2 d3a3bf1 37abfc2 38d816c d3a3bf1 37abfc2 8124057 37abfc2 8124057 37abfc2 38d816c 8124057 38d816c 37abfc2 38d816c 37abfc2 fdd3761 38d816c 8124057 38d816c fdd3761 37abfc2 fdd3761 37abfc2 38d816c 37abfc2 4e7c237 38d816c 8124057 37abfc2 38d816c b2ef34b |
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 |
# app.py
import gradio as gr
from PIL import Image
import os
import torch
from diffusers import StableDiffusionInpaintPipeline
from utils.planner import extract_scene_plan # 🧠 Brain Layer
# ----------------------------
# 🔧 Device Setup
# ----------------------------
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.float16 if device == "cuda" else torch.float32
# ----------------------------
# 📦 Load Inpainting Model
# ----------------------------
pipe = StableDiffusionInpaintPipeline.from_pretrained(
"runwayml/stable-diffusion-inpainting",
torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
).to("cuda" if torch.cuda.is_available() else "cpu")
pipe.enable_attention_slicing()
pipe.enable_model_cpu_offload()
# ----------------------------
# 🎨 Image Generation Function
# ----------------------------
def process_image(prompt, image, mask, num_variations):
try:
print("🧠 Prompt received:", prompt)
# 🧠 Step 1: Brain Layer
reasoning_json = extract_scene_plan(prompt)
print("🧠 Scene plan extracted:", reasoning_json)
# Resize inputs to 1024x1024 (required for SDXL)
image = image.resize((1024, 1024)).convert("RGB")
mask = mask.resize((1024, 1024)).convert("L")
results = []
for i in range(num_variations):
print(f"🎨 Generating variation {i + 1}...")
output = pipe(
prompt=prompt,
image=image,
mask_image=mask,
strength=0.98,
guidance_scale=7.5,
num_inference_steps=40
).images[0]
results.append(output)
return results, reasoning_json
except Exception as e:
print("❌ Error during generation:", e)
return ["❌ Generation failed"], {"error": str(e)}
# ----------------------------
# 🖼️ Gradio UI
# ----------------------------
with gr.Blocks() as demo:
gr.Markdown("## 🧠 NewCrux Inpainting Demo (SDXL)\nUpload a product image, a mask, and a prompt to generate realistic content.")
with gr.Row():
with gr.Column():
prompt_input = gr.Textbox(label="Prompt")
image_input = gr.Image(type="pil", label="Upload Product Image")
mask_input = gr.Image(type="pil", label="Upload Mask (white = keep, black = replace)")
variation_slider = gr.Slider(1, 4, step=1, value=1, label="Number of Variations")
generate_btn = gr.Button("Generate")
with gr.Column():
output_gallery = gr.Gallery(
label="Generated Variations",
columns=2,
rows=2,
height="auto"
)
json_output = gr.JSON(label="🧠 Brain Layer Reasoning")
generate_btn.click(
fn=process_image,
inputs=[prompt_input, image_input, mask_input, variation_slider],
outputs=[output_gallery, json_output]
)
demo.launch()
|