File size: 5,458 Bytes
7941d68 | 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 158 159 160 161 162 | import gradio as gr
import spaces
import torch
from diffusers import DiffusionPipeline
import numpy as np
from PIL import Image
import os
from rembg import remove
# --- MODELS LOADING (Hugging Face Cache se load honge) ---
print("Loading FLUX (Text-to-Image) model...")
# FLUX.1-schnell use kar rahe hain speed ke liye
txt2img_pipe = DiffusionPipeline.from_pretrained(
"black-forest-labs/FLUX.1-schnell",
torch_dtype=torch_dtype,
device_map="auto"
)
print("Loading Image-to-3D model...")
# Hum standard shape generation pipeline loading import use karenge
# (Hugging Face par different Space backends TripoSR ko different tarah se import karte hain,
# free Tier ke liye yeh standard approach hai)
try:
from tsr.system import TSR
from tsr.utils import remove_background, resize_foreground, save_video
IS_TSR_AVAILABLE = True
# TSR Model Load (CPU par initialize, GPU decorator switch karega)
tsr_model = TSR.from_pretrained(
"stabilityai/TripoSR",
config_name="config.yaml",
weight_name="model.ckpt"
)
except ImportError:
print("Warning: TripoSR libraries not found. Image-to-3D will be dummy.")
IS_TSR_AVAILABLE = False
# --- HELPERS ---
def process_image_background(image_path):
"""Image se background hatata hai 3D generation se pehle"""
input_image = Image.open(image_path)
# Using rembg for background removal
output_image = remove(input_image)
# Save as PNG with alpha channel
bg_removed_path = "processed_input.png"
output_image.save(bg_removed_path)
return bg_removed_path
def generate_3d_from_processed_image(processed_image_path):
"""TripoSR ka use karke PNG se .glb banata hai"""
if not IS_TSR_AVAILABLE:
return None # Backend error safety
# Run TripoSR generation
# ZeroGPU handle karega move to device automatically decorator ke sath
scene_codes = tsr_model(processed_image_path, device="cuda")
# Mesh extract karke .glb format mein save karna
# (Extracting vertices/faces internally and saving)
meshes = tsr_model.extract_mesh(scene_codes)
glb_path = "generated_model.glb"
meshes[0].export(glb_path) # Export first generated mesh
return glb_path
# --- GENERATION FUNCTIONS (with GPU Decorator) ---
@spaces.GPU(duration=60) # Text to Image thoda time leta hai
def text_to_3d_pipeline(prompt):
if not prompt:
return None, None
print(f"Generating Image for prompt: {prompt}")
# 1. Text to Real Image (FLUX)
# inference_steps kam rakhe hain speed ke liye
image_result = txt2img_pipe(
prompt,
guidance_scale=0.0,
num_inference_steps=4,
max_sequence_length=256
).images[0]
temp_img_path = "text2img_output.png"
image_result.save(temp_img_path)
print("Generating 3D model from image...")
# 2. Background Removal & 3D Generation
if IS_TSR_AVAILABLE:
processed_img = process_image_background(temp_img_path)
glb_file = generate_3d_from_processed_image(processed_img)
return temp_img_path, glb_file
else:
return temp_img_path, None # Image dikhegi par 3D nahi banega safety ke liye
@spaces.GPU(duration=30) # Image to 3D faster hota hai
def image_to_3d_pipeline(image_filepath):
if not image_filepath or not IS_TSR_AVAILABLE:
return None
print("Processing uploaded image for 3D...")
# 1. Background Removal
processed_img = process_image_background(image_filepath)
# 2. 3D Generation
glb_file = generate_3d_from_processed_image(processed_img)
return glb_file
# --- UI DESIGN ---
with gr.Blocks(theme=gr.themes.Soft()) as demo:
gr.Markdown("# 🌌 V.O.I.D - Real 3D Generation Studio")
gr.Markdown("Generating real images and 3D models using FLUX.1 & TripoSR (ZeroGPU).")
with gr.Row():
# Left Side: Input Options
with gr.Column(scale=1):
with gr.Tabs():
# Tab 1: Text to 3D
with gr.TabItem("📝 Text to 3D Pipeline"):
txt_prompt = gr.Textbox(
label="Describe your 3D model",
placeholder="e.g., A photorealistic golden crown, a cyberpunk helmet..."
)
txt_btn = gr.Button("Generate Image & 3D 🚀", variant="primary")
# Intermediate Image Output
gen_img_output = gr.Image(label="Generated Reference Image", type="filepath")
# Tab 2: Image to 3D
with gr.TabItem("🖼️ Image to 3D"):
img_input = gr.Image(type="filepath", label="Upload Photo (Background will be removed)")
img_btn = gr.Button("Generate 3D Model 🚀", variant="primary")
# Right Side: Final 3D Model Viewer
with gr.Column(scale=1):
model_viewer = gr.Model3D(
label="Final 3D Model Viewer (.glb)",
height=500
)
# --- LOGIC CONNECTIONS ---
# Text-to-3D returns [Image, GLB]
txt_btn.click(
fn=text_to_3d_pipeline,
inputs=[txt_prompt],
outputs=[gen_img_output, model_viewer]
)
# Image-to-3D returns [GLB]
img_btn.click(
fn=image_to_3d_pipeline,
inputs=[img_input],
outputs=[model_viewer]
)
# Launch
demo.launch()
|