lora

app.py

@@ -1,27 +1,20 @@
 import spaces
-import argparse
 import os
 import time
 from os import path
-from safetensors.torch import load_file
 from huggingface_hub import hf_hub_download
-import imageio
 import numpy as np
 import torch
 import rembg
 from PIL import Image
 from torchvision.transforms import v2
+from einops import rearrange
 from pytorch_lightning import seed_everything
 from omegaconf import OmegaConf
-from einops import rearrange, repeat
-from tqdm import tqdm
 from diffusers import DiffusionPipeline, EulerAncestralDiscreteScheduler
 import gradio as gr
 import shutil
 import tempfile
-from functools import partial
-from optimum.quanto import quantize, qfloat8, freeze
-from diffusers import FluxPipeline
 from src.utils.train_util import instantiate_from_config
 from src.utils.camera_util import (
     FOV_to_intrinsics, 
@@ -30,6 +23,9 @@ from src.utils.camera_util import (
 )
 from src.utils.mesh_util import save_obj, save_glb
 from src.utils.infer_util import remove_background, resize_foreground, images_to_video
+import random
+import requests
+import io
 
 # Set up cache path
 cache_path = path.join(path.dirname(path.abspath(__file__)), "models")
@@ -71,19 +67,11 @@ else:
     print("CUDA installation not found")
 
 
-device = 'cuda'
-
-base_model = "black-forest-labs/FLUX.1-dev"
-file_flux = hf_hub_download("marduk191/Flux.1_collection", "flux.1_dev_8x8_e4m3fn-marduk191.safetensors")
-pipe = FluxPipeline.from_single_file(file_flux, torch_dtype=torch.bfloat16, token=huggingface_token)
-
-# Load and fuse LoRA BEFORE quantizing
-print('Loading and fusing lora, please wait...')
-lora_path = hf_hub_download("gokaygokay/Flux-Game-Assets-LoRA-v2", "game_asst.safetensors")
-pipe.load_lora_weights(lora_path)
-pipe.fuse_lora(lora_scale=1.0)
-pipe.unload_lora_weights()
+API_TOKEN = os.getenv("HUGGINGFACE_TOKEN")
+headers = {"Authorization": f"Bearer {API_TOKEN}"}
+timeout = 100
 
+device = 'cuda'
 
 # Load 3D generation models
 config_path = 'configs/instant-mesh-large.yaml'
@@ -143,20 +131,7 @@ def preprocess(input_image, do_remove_background):
         input_image = resize_foreground(input_image, 0.85)
     return input_image
 
-ts_cutoff = 2
 
-@spaces.GPU
-def generate_flux_image(prompt, height, width, steps, scales, seed):
-    pipe.to(device)
-    return pipe(
-        prompt=prompt, 
-        width=int(height),
-        height=int(width),
-        num_inference_steps=int(steps), 
-        generator=torch.Generator().manual_seed(int(seed)),
-        guidance_scale=float(scales),
-        timestep_to_start_cfg=ts_cutoff,
-    ).images[0]
 
 
 @spaces.GPU
@@ -209,6 +184,45 @@ def make3d(images):
     
     return mesh_fpath, mesh_glb_fpath
 
+# Remove the FluxPipeline setup and replace with the query function
+def query(prompt, steps=28, cfg_scale=3.5, randomize_seed=True, seed=-1, width=1024, height=1024):
+    if not prompt:
+        return None
+
+    lora_id = "gokaygokay/Flux-Game-Assets-LoRA-v2"
+    API_URL = f"https://api-inference.huggingface.co/models/{lora_id}"
+    
+    if randomize_seed:
+        seed = random.randint(1, 4294967296)
+    
+    prompt = f"{prompt} | ultra detail, ultra elaboration, ultra quality, perfect."
+    
+    payload = {
+        "inputs": prompt,
+        "steps": steps,
+        "cfg_scale": cfg_scale,
+        "seed": seed,
+        "parameters": {
+            "width": width,
+            "height": height
+        }
+    }
+
+    response = requests.post(API_URL, headers=headers, json=payload, timeout=100)
+    if response.status_code != 200:
+        if response.status_code == 503:
+            raise gr.Error("The model is being loaded")
+        raise gr.Error(f"Error {response.status_code}")
+    
+    try:
+        image_bytes = response.content
+        image = Image.open(io.BytesIO(image_bytes))
+        return image
+    except Exception as e:
+        print(f"Error when trying to open the image: {e}")
+        return None
+
+# Update the Gradio interface
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown(
         """
@@ -236,7 +250,8 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
                         steps = gr.Slider(label="Inference Steps", minimum=10, maximum=50, step=1, value=28)
                         scales = gr.Slider(label="Guidance Scale", minimum=0.0, maximum=5.0, step=0.1, value=3.5)
                     
-                    seed = gr.Number(label="Seed (for reproducibility)", value=3413, precision=0)
+                    seed = gr.Number(label="Seed", value=-1, precision=0)
+                    randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
             
             generate_btn = gr.Button("Generate 3D Model", variant="primary")
 
@@ -251,8 +266,13 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
 
     mv_images = gr.State()
 
-    def process_pipeline(prompt, height, width, steps, scales, seed):
-        flux_image = generate_flux_image(prompt, height, width, steps, scales, seed)
+    def process_pipeline(prompt, height, width, steps, scales, seed, randomize_seed):
+        # Generate Flux image using the API
+        prompt_real = f"wbgmsst, {prompt}, white background"
+        flux_image = query(prompt_real, steps, scales, randomize_seed, seed, width, height)
+        if flux_image is None:
+            raise gr.Error("Failed to generate image")
+            
         processed_image = preprocess(flux_image, do_remove_background=True)
         mv_images, show_image = generate_mvs(processed_image, steps, seed)
         obj_path, glb_path = make3d(mv_images)
@@ -260,7 +280,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
 
     generate_btn.click(
         fn=process_pipeline,
-        inputs=[prompt, height, width, steps, scales, seed],
+        inputs=[prompt, height, width, steps, scales, seed, randomize_seed],
         outputs=[flux_output, mv_show_images, output_model_obj, output_model_glb]
     )