app.py

# import gradio as gr
# import torch
# from PIL import Image
# from model import CRM
# from inference import generate3d
# import numpy as np

# # Load model
# crm_path = "CRM.pth"  # Make sure the model is uploaded to the Space
# model = CRM(torch.load(crm_path, map_location="cpu"))
# model = model.to("cuda:0" if torch.cuda.is_available() else "cpu")

# def generate_3d(image_path, seed=1234, scale=5.5, step=30):
#     image = Image.open(image_path).convert("RGB")
#     np_img = np.array(image)
#     glb_path = generate3d(model, np_img, np_img, "cuda:0" if torch.cuda.is_available() else "cpu")
#     return glb_path

# iface = gr.Interface(
#     fn=generate_3d,
#     inputs=gr.Image(type="filepath"),
#     outputs=gr.Model3D(),
#     title="Convolutional Reconstruction Model (CRM)",
#     description="Upload an image to generate a 3D model."
# )

# iface.launch()


#############2nd################3
# import os
# import torch
# import gradio as gr
# from huggingface_hub import hf_hub_download
# from model import CRM  # Make sure this matches your model file structure

# # Define model details
# REPO_ID = "Mariam-Elz/CRM"  # Hugging Face model repo
# MODEL_FILES = {
#     "ccm-diffusion": "ccm-diffusion.pth",
#     "pixel-diffusion": "pixel-diffusion.pth",
#     "CRM": "CRM.pth"
# }
# DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# # Download models from Hugging Face if not already present
# MODEL_DIR = "./models"
# os.makedirs(MODEL_DIR, exist_ok=True)

# for name, filename in MODEL_FILES.items():
#     model_path = os.path.join(MODEL_DIR, filename)
#     if not os.path.exists(model_path):
#         print(f"Downloading {filename}...")
#         hf_hub_download(repo_id=REPO_ID, filename=filename, local_dir=MODEL_DIR)

# # Load the model
# print("Loading CRM Model...")
# model = CRM()
# model.load_state_dict(torch.load(os.path.join(MODEL_DIR, MODEL_FILES["CRM"]), map_location=DEVICE))
# model.to(DEVICE)
# model.eval()
# print("✅ Model Loaded Successfully!")

# # Define Gradio Interface
# def predict(input_image):
#     with torch.no_grad():
#         output = model(input_image.to(DEVICE))  # Modify based on model input format
#     return output.cpu()

# demo = gr.Interface(
#     fn=predict,
#     inputs=gr.Image(type="pil"),
#     outputs=gr.Image(type="pil"),
#     title="Convolutional Reconstruction Model (CRM)",
#     description="Upload an image to generate a reconstructed output."
# )

# if __name__ == "__main__":
#     demo.launch()
########################3rd-MAIN######################3

# import torch
# import gradio as gr
# import requests
# import os

# # Download model weights from Hugging Face model repo (if not already present)
# model_repo = "Mariam-Elz/CRM"  # Your Hugging Face model repo

# model_files = {
#     "ccm-diffusion.pth": "ccm-diffusion.pth",
#     "pixel-diffusion.pth": "pixel-diffusion.pth",
#     "CRM.pth": "CRM.pth",
# }

# os.makedirs("models", exist_ok=True)

# for filename, output_path in model_files.items():
#     file_path = f"models/{output_path}"
#     if not os.path.exists(file_path):
#         url = f"https://huggingface.co/{model_repo}/resolve/main/{filename}"
#         print(f"Downloading {filename}...")
#         response = requests.get(url)
#         with open(file_path, "wb") as f:
#             f.write(response.content)

# # Load model (This part depends on how the model is defined)
# device = "cuda" if torch.cuda.is_available() else "cpu"

# def load_model():
#     model_path = "models/CRM.pth"
#     model = torch.load(model_path, map_location=device)
#     model.eval()
#     return model

# model = load_model()

# # Define inference function
# def infer(image):
#     """Process input image and return a reconstructed image."""
#     with torch.no_grad():
#         # Assuming model expects a tensor input
#         image_tensor = torch.tensor(image).to(device)
#         output = model(image_tensor)
#         return output.cpu().numpy()

# # Create Gradio UI
# demo = gr.Interface(
#     fn=infer,
#     inputs=gr.Image(type="numpy"),
#     outputs=gr.Image(type="numpy"),
#     title="Convolutional Reconstruction Model",
#     description="Upload an image to get the reconstructed output."
# )

# if __name__ == "__main__":
#     demo.launch()


#################4th##################

# import torch
# import gradio as gr
# import requests
# import os

# # Define model repo
# model_repo = "Mariam-Elz/CRM"

# # Define model files and download paths
# model_files = {
#     "CRM.pth": "models/CRM.pth"
# }

# os.makedirs("models", exist_ok=True)

# # Download model files only if they don't exist
# for filename, output_path in model_files.items():
#     if not os.path.exists(output_path):
#         url = f"https://huggingface.co/{model_repo}/resolve/main/{filename}"
#         print(f"Downloading {filename}...")
#         response = requests.get(url)
#         with open(output_path, "wb") as f:
#             f.write(response.content)

# # Load model with low memory usage
# def load_model():
#     model_path = "models/CRM.pth"
#     model = torch.load(model_path, map_location="cpu")  # Load on CPU to reduce memory usage
#     model.eval()
#     return model

# model = load_model()

# # Define inference function
# def infer(image):
#     """Process input image and return a reconstructed image."""
#     with torch.no_grad():
#         image_tensor = torch.tensor(image).unsqueeze(0)  # Add batch dimension
#         image_tensor = image_tensor.to("cpu")  # Keep on CPU to save memory
#         output = model(image_tensor)
#         return output.squeeze(0).numpy()

# # Create Gradio UI
# demo = gr.Interface(
#     fn=infer,
#     inputs=gr.Image(type="numpy"),
#     outputs=gr.Image(type="numpy"),
#     title="Convolutional Reconstruction Model",
#     description="Upload an image to get the reconstructed output."
# )

# if __name__ == "__main__":
#     demo.launch()


# ##############5TH#################
# import torch
# import torch.nn as nn
# import gradio as gr
# import requests
# import os

# # Define model repo
# model_repo = "Mariam-Elz/CRM"

# # Define model files and download paths
# model_files = {
#     "CRM.pth": "models/CRM.pth"
# }

# os.makedirs("models", exist_ok=True)

# # Download model files only if they don't exist
# for filename, output_path in model_files.items():
#     if not os.path.exists(output_path):
#         url = f"https://huggingface.co/{model_repo}/resolve/main/{filename}"
#         print(f"Downloading {filename}...")
#         response = requests.get(url)
#         with open(output_path, "wb") as f:
#             f.write(response.content)

# # Define the model architecture (you MUST replace this with your actual model)
# class CRM_Model(nn.Module):
#     def __init__(self):
#         super(CRM_Model, self).__init__()
#         self.layer1 = nn.Conv2d(3, 64, kernel_size=3, padding=1)
#         self.relu = nn.ReLU()
#         self.layer2 = nn.Conv2d(64, 3, kernel_size=3, padding=1)

#     def forward(self, x):
#         x = self.layer1(x)
#         x = self.relu(x)
#         x = self.layer2(x)
#         return x

# # Load model with proper architecture
# def load_model():
#     model = CRM_Model()  # Instantiate the model architecture
#     model_path = "models/CRM.pth"
#     model.load_state_dict(torch.load(model_path, map_location="cpu"))  # Load weights
#     model.eval()  # Set to evaluation mode
#     return model

# model = load_model()

# # Define inference function
# def infer(image):
#     """Process input image and return a reconstructed image."""
#     with torch.no_grad():
#         image_tensor = torch.tensor(image).unsqueeze(0).permute(0, 3, 1, 2).float() / 255.0  # Convert to tensor
#         output = model(image_tensor)  # Run through model
#         output = output.squeeze(0).permute(1, 2, 0).numpy() * 255.0  # Convert back to image
#         return output.astype("uint8")

# # Create Gradio UI
# demo = gr.Interface(
#     fn=infer,
#     inputs=gr.Image(type="numpy"),
#     outputs=gr.Image(type="numpy"),
#     title="Convolutional Reconstruction Model",
#     description="Upload an image to get the reconstructed output."
# )

# if __name__ == "__main__":
#     demo.launch()


#############6th-worked-proc##################
# import torch
# import gradio as gr
# import requests
# import os
# import numpy as np

# # Hugging Face Model Repository
# model_repo = "Mariam-Elz/CRM"

# # Download Model Weights (Only CRM.pth to Save Memory)
# model_path = "models/CRM.pth"
# os.makedirs("models", exist_ok=True)

# if not os.path.exists(model_path):
#     url = f"https://huggingface.co/{model_repo}/resolve/main/CRM.pth"
#     print(f"Downloading CRM.pth...")
#     response = requests.get(url)
#     with open(model_path, "wb") as f:
#         f.write(response.content)

# # Set Device (Use CPU to Reduce RAM Usage)
# device = "cpu"

# # Load Model Efficiently
# def load_model():
#     model = torch.load(model_path, map_location=device)
#     if isinstance(model, torch.nn.Module):
#         model.eval()  # Ensure model is in inference mode
#     return model

# # Load model only when needed (saves memory)
# model = load_model()

# # Define Inference Function with Memory Optimizations
# def infer(image):
#     """Process input image and return a reconstructed image."""
#     with torch.no_grad():
#         # Convert image to torch tensor & normalize (float16 to save RAM)
#         image_tensor = torch.tensor(image, dtype=torch.float16).unsqueeze(0).permute(0, 3, 1, 2) / 255.0
#         image_tensor = image_tensor.to(device)

#         # Model Inference
#         output = model(image_tensor)

#         # Convert back to numpy image format
#         output_image = output.squeeze(0).permute(1, 2, 0).cpu().numpy() * 255.0
#         output_image = np.clip(output_image, 0, 255).astype(np.uint8)

#         # Free Memory
#         del image_tensor, output
#         torch.cuda.empty_cache()

#         return output_image

# # Create Gradio UI
# demo = gr.Interface(
#     fn=infer,
#     inputs=gr.Image(type="numpy"),
#     outputs=gr.Image(type="numpy"),
#     title="Optimized Convolutional Reconstruction Model",
#     description="Upload an image to get the reconstructed output with reduced memory usage."
# )

# if __name__ == "__main__":
#     demo.launch()



#############7tth################
# import torch
# import torch.nn as nn
# import gradio as gr
# import requests
# import os
# import torchvision.transforms as transforms
# import numpy as np
# from PIL import Image

# # Hugging Face Model Repository
# model_repo = "Mariam-Elz/CRM"

# # Model File Path
# model_path = "models/CRM.pth"
# os.makedirs("models", exist_ok=True)

# # Download model weights if not present
# if not os.path.exists(model_path):
#     url = f"https://huggingface.co/{model_repo}/resolve/main/CRM.pth"
#     print(f"Downloading CRM.pth...")
#     response = requests.get(url)
#     with open(model_path, "wb") as f:
#         f.write(response.content)

# # Set Device
# device = "cuda" if torch.cuda.is_available() else "cpu"

# # Define Model Architecture (Replace with your actual model)
# class CRMModel(nn.Module):
#     def __init__(self):
#         super(CRMModel, self).__init__()
#         self.conv1 = nn.Conv2d(3, 64, kernel_size=3, padding=1)
#         self.conv2 = nn.Conv2d(64, 64, kernel_size=3, padding=1)
#         self.relu = nn.ReLU()
    
#     def forward(self, x):
#         x = self.relu(self.conv1(x))
#         x = self.relu(self.conv2(x))
#         return x

# # Load Model
# def load_model():
#     print("Loading model...")
#     model = CRMModel()  # Use the correct architecture here
#     state_dict = torch.load(model_path, map_location=device)

#     if isinstance(state_dict, dict):  # Ensure it's a valid state_dict
#         model.load_state_dict(state_dict)
#     else:
#         raise ValueError("Error: The loaded state_dict is not in the correct format.")

#     model.to(device)
#     model.eval()
#     print("Model loaded successfully!")
#     return model

# # Load the model
# model = load_model()

# # Define Inference Function
# def infer(image):
#     """Process input image and return a reconstructed 3D output."""
#     try:
#         print("Preprocessing image...")

#         # Convert image to PyTorch tensor & normalize
#         transform = transforms.Compose([
#             transforms.Resize((256, 256)),  # Resize to fit model input
#             transforms.ToTensor(),  # Converts to tensor (C, H, W)
#             transforms.Normalize(mean=[0.5], std=[0.5]),  # Normalize
#         ])
#         image_tensor = transform(image).unsqueeze(0).to(device)  # Add batch dimension

#         print("Running inference...")
#         with torch.no_grad():
#             output = model(image_tensor)  # Forward pass

#         # Ensure output is a valid tensor
#         if isinstance(output, torch.Tensor):
#             output_image = output.squeeze(0).permute(1, 2, 0).cpu().numpy()
#             output_image = np.clip(output_image * 255.0, 0, 255).astype(np.uint8)
#             print("Inference complete! Returning output.")
#             return output_image
#         else:
#             print("Error: Model output is not a tensor.")
#             return None

#     except Exception as e:
#         print(f"Error during inference: {e}")
#         return None

# # Create Gradio UI
# demo = gr.Interface(
#     fn=infer,
#     inputs=gr.Image(type="pil"),
#     outputs=gr.Image(type="numpy"),
#     title="Convolutional Reconstruction Model",
#     description="Upload an image to get the reconstructed output."
# )

# if __name__ == "__main__":
#     demo.launch()




# Not ready to use yet
import spaces
import argparse
import numpy as np
import gradio as gr
from omegaconf import OmegaConf
import torch
from PIL import Image
import PIL
from pipelines import TwoStagePipeline
from huggingface_hub import hf_hub_download
import os
import rembg
from typing import Any
import json
import os
import json
import argparse

from model import CRM
from inference import generate3d

pipeline = None
rembg_session = rembg.new_session()


def expand_to_square(image, bg_color=(0, 0, 0, 0)):
    # expand image to 1:1
    width, height = image.size
    if width == height:
        return image
    new_size = (max(width, height), max(width, height))
    new_image = Image.new("RGBA", new_size, bg_color)
    paste_position = ((new_size[0] - width) // 2, (new_size[1] - height) // 2)
    new_image.paste(image, paste_position)
    return new_image

def check_input_image(input_image):
    if input_image is None:
        raise gr.Error("No image uploaded!")


def remove_background(
    image: PIL.Image.Image,
    rembg_session: Any = None,
    force: bool = False,
    **rembg_kwargs,
) -> PIL.Image.Image:
    do_remove = True
    if image.mode == "RGBA" and image.getextrema()[3][0] < 255:
        # explain why current do not rm bg
        print("alhpa channl not enpty, skip remove background, using alpha channel as mask")
        background = Image.new("RGBA", image.size, (0, 0, 0, 0))
        image = Image.alpha_composite(background, image)
        do_remove = False
    do_remove = do_remove or force
    if do_remove:
        image = rembg.remove(image, session=rembg_session, **rembg_kwargs)
    return image

def do_resize_content(original_image: Image, scale_rate):
    # resize image content wile retain the original image size
    if scale_rate != 1:
        # Calculate the new size after rescaling
        new_size = tuple(int(dim * scale_rate) for dim in original_image.size)
        # Resize the image while maintaining the aspect ratio
        resized_image = original_image.resize(new_size)
        # Create a new image with the original size and black background
        padded_image = Image.new("RGBA", original_image.size, (0, 0, 0, 0))
        paste_position = ((original_image.width - resized_image.width) // 2, (original_image.height - resized_image.height) // 2)
        padded_image.paste(resized_image, paste_position)
        return padded_image
    else:
        return original_image

def add_background(image, bg_color=(255, 255, 255)):
    # given an RGBA image, alpha channel is used as mask to add background color
    background = Image.new("RGBA", image.size, bg_color)
    return Image.alpha_composite(background, image)


def preprocess_image(image, background_choice, foreground_ratio, backgroud_color):
    """
    input image is a pil image in RGBA, return RGB image
    """
    print(background_choice)
    if background_choice == "Alpha as mask":
        background = Image.new("RGBA", image.size, (0, 0, 0, 0))
        image = Image.alpha_composite(background, image)
    else:
        image = remove_background(image, rembg_session, force=True)
    image = do_resize_content(image, foreground_ratio)
    image = expand_to_square(image)
    image = add_background(image, backgroud_color)
    return image.convert("RGB")

@spaces.GPU
def gen_image(input_image, seed, scale, step):
    global pipeline, model, args
    pipeline.set_seed(seed)
    rt_dict = pipeline(input_image, scale=scale, step=step)
    stage1_images = rt_dict["stage1_images"]
    stage2_images = rt_dict["stage2_images"]
    np_imgs = np.concatenate(stage1_images, 1)
    np_xyzs = np.concatenate(stage2_images, 1)

    glb_path = generate3d(model, np_imgs, np_xyzs, args.device)
    return Image.fromarray(np_imgs), Image.fromarray(np_xyzs), glb_path#, obj_path


parser = argparse.ArgumentParser()
parser.add_argument(
    "--stage1_config",
    type=str,
    default="configs/nf7_v3_SNR_rd_size_stroke.yaml",
    help="config for stage1",
)
parser.add_argument(
    "--stage2_config",
    type=str,
    default="configs/stage2-v2-snr.yaml",
    help="config for stage2",
)

parser.add_argument("--device", type=str, default="cuda")
args = parser.parse_args()

crm_path = hf_hub_download(repo_id="Zhengyi/CRM", filename="CRM.pth")
specs = json.load(open("configs/specs_objaverse_total.json"))
model = CRM(specs)
model.load_state_dict(torch.load(crm_path, map_location="cpu"), strict=False)
model = model.to(args.device)

stage1_config = OmegaConf.load(args.stage1_config).config
stage2_config = OmegaConf.load(args.stage2_config).config
stage2_sampler_config = stage2_config.sampler
stage1_sampler_config = stage1_config.sampler

stage1_model_config = stage1_config.models
stage2_model_config = stage2_config.models

xyz_path = hf_hub_download(repo_id="Zhengyi/CRM", filename="ccm-diffusion.pth")
pixel_path = hf_hub_download(repo_id="Zhengyi/CRM", filename="pixel-diffusion.pth")
stage1_model_config.resume = pixel_path
stage2_model_config.resume = xyz_path

pipeline = TwoStagePipeline(
    stage1_model_config,
    stage2_model_config,
    stage1_sampler_config,
    stage2_sampler_config,
    device=args.device,
    dtype=torch.float32
)

_DESCRIPTION = '''
* Our [official implementation](https://github.com/thu-ml/CRM) uses UV texture instead of vertex color. It has better texture than this online demo.
* Project page of CRM: https://ml.cs.tsinghua.edu.cn/~zhengyi/CRM/
* If you find the output unsatisfying, try using different seeds:)
'''

with gr.Blocks() as demo:
    gr.Markdown("# CRM: Single Image to 3D Textured Mesh with Convolutional Reconstruction Model")
    gr.Markdown(_DESCRIPTION)
    with gr.Row():
        with gr.Column():
            with gr.Row():
                image_input = gr.Image(
                    label="Image input",
                    image_mode="RGBA",
                    sources="upload",
                    type="pil",
                )
                processed_image = gr.Image(label="Processed Image", interactive=False, type="pil", image_mode="RGB")
            with gr.Row():
                with gr.Column():
                    with gr.Row():
                        background_choice = gr.Radio([
                                "Alpha as mask",
                                "Auto Remove background"
                            ], value="Auto Remove background",
                            label="backgroud choice")
                        # do_remove_background = gr.Checkbox(label=, value=True)
                        # force_remove = gr.Checkbox(label=, value=False)
                    back_groud_color = gr.ColorPicker(label="Background Color", value="#7F7F7F", interactive=False)
                    foreground_ratio = gr.Slider(
                        label="Foreground Ratio",
                        minimum=0.5,
                        maximum=1.0,
                        value=1.0,
                        step=0.05,
                    )

                with gr.Column():
                    seed = gr.Number(value=1234, label="seed", precision=0)
                    guidance_scale = gr.Number(value=5.5, minimum=3, maximum=10, label="guidance_scale")
                    step = gr.Number(value=30, minimum=30, maximum=100, label="sample steps", precision=0)
            text_button = gr.Button("Generate 3D shape")
            gr.Examples(
                examples=[os.path.join("examples", i) for i in os.listdir("examples")],
                inputs=[image_input],
                examples_per_page = 20,
            )
        with gr.Column():
            image_output = gr.Image(interactive=False, label="Output RGB image")
            xyz_ouput = gr.Image(interactive=False, label="Output CCM image")

            output_model = gr.Model3D(
                label="Output OBJ",
                interactive=False,
            )
            gr.Markdown("Note: Ensure that the input image is correctly pre-processed into a grey background, otherwise the results will be unpredictable.")

    inputs = [
        processed_image,
        seed,
        guidance_scale,
        step,
    ]
    outputs = [
        image_output,
        xyz_ouput,
        output_model,
        # output_obj,
    ]


    text_button.click(fn=check_input_image, inputs=[image_input]).success(
        fn=preprocess_image,
        inputs=[image_input, background_choice, foreground_ratio, back_groud_color],
        outputs=[processed_image],
    ).success(
        fn=gen_image,
        inputs=inputs,
        outputs=outputs,
    )
    demo.queue().launch()