app.py

import gradio as gr
import torch
import gc
import os
import numpy as np
from PIL import Image
from diffusers import StableDiffusionPipeline

# 3D 관련 라이브러리 (Lazy Import to save memory)
from tsr.system import TSR
from tsr.utils import remove_background, resize_foreground, save_video

# CPU 설정
device = "cpu"

def generate_pipeline(prompt):
    output_obj_path = "output_model.obj"
    
    # ==========================================
    # 1단계: Text > 2D 이미지 생성 (IntimeAI/Miro)
    # ==========================================
    print(f">>> [Step 1/3] Loading 2D Model (Miro) for: {prompt}")
    
    # 2D 모델 로드
    pipe = StableDiffusionPipeline.from_pretrained(
        "IntimeAI/Miro", 
        torch_dtype=torch.float32,
        use_safetensors=True
    ).to(device)
    
    # 이미지 생성
    image = pipe(prompt, num_inference_steps=20).images[0]
    
    # 이미지 저장 (디버깅용)
    image_path = "generated_2d.png"
    image.save(image_path)
    print(">>> 2D Image Generated!")

    # 중요: 메모리 확보를 위해 2D 모델 삭제 (무료 서버 필수)
    del pipe
    gc.collect()
    print(">>> 2D Model unloaded to free up RAM.")

    # ==========================================
    # 2단계: 이미지 전처리 (배경 제거)
    # ==========================================
    print(f">>> [Step 2/3] Removing Background...")
    
    # 배경 제거 (rembg)
    rembg_session = None # 세션 자동 생성
    image_no_bg = remove_background(image, rembg_session)
    image_no_bg = resize_foreground(image_no_bg, 0.85)
    
    # ==========================================
    # 3단계: 2D > 3D 모델 생성 (TripoSR)
    # ==========================================
    print(f">>> [Step 3/3] Loading 3D Model (TripoSR)...")
    
    # 3D 모델 로드
    model_3d = TSR.from_pretrained(
        "stabilityai/TripoSR",
        config_name="config.yaml",
        weight_name="model.ckpt",
    )
    model_3d.to(device)
    
    # 3D 생성
    scene_codes = model_3d(image_no_bg, device=device)
    meshes = model_3d.extract_mesh(scene_codes)[0]
    
    # OBJ 저장
    meshes.export(output_obj_path)
    print(">>> 3D OBJ Saved!")

    # 3D 모델 삭제
    del model_3d
    gc.collect()

    return output_obj_path

# Gradio 인터페이스
iface = gr.Interface(
    fn=generate_pipeline,
    inputs=gr.Textbox(label="Enter Prompt (e.g., A isometric sword in fantasy style)"),
    outputs=gr.Model3D(label="Generated 3D Object", clear_color=[0, 0, 0, 0]),
    title="Miro to 3D Generator (Text -> 2D -> 3D)",
    description="IntimeAI/Miro 모델을 사용하여 이미지를 만들고 3D로 변환합니다. (CPU 모드: 약 5~10분 소요)"
)

iface.launch()