Spaces:

Tejasf
/

Newera

Runtime error

Newera / app.py

Create app.py

7941d68 verified 4 months ago

5.46 kB

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()

	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()