Refactor code for improved performance and readability

.gitignore

@@ -0,0 +1,2 @@
+.env
+output.png

TripoSR

@@ -0,0 +1 @@
+Subproject commit 8e51fec8095c9eae20e6ea7c9aef6368c5631a21

app.py

@@ -0,0 +1,13 @@
+import gradio as gr
+from worker import worker  # import the worker function
+
+def greet(name):
+    return "Hello " + name + "!!"
+
+def kickoff_worker():
+    worker()  # call the worker function
+
+iface = gr.Interface(fn=greet, inputs="text", outputs="text")
+iface.launch()
+print("Launching worker...")
+kickoff_worker()  # kickoff the worker after launching the interface

imgGen.py

@@ -0,0 +1,28 @@
+import torch
+from diffusers import StableDiffusionXLPipeline, UNet2DConditionModel, EulerDiscreteScheduler
+from huggingface_hub import hf_hub_download
+from safetensors.torch import load_file
+
+
+base = "stabilityai/stable-diffusion-xl-base-1.0"
+repo = "ByteDance/SDXL-Lightning"
+ckpt = "sdxl_lightning_4step_unet.safetensors" # Use the correct ckpt for your step setting!
+
+
+# Load model.
+unet = UNet2DConditionModel.from_config(base, subfolder="unet").to("cuda", torch.float16)
+unet.load_state_dict(load_file(hf_hub_download(repo, ckpt), device="cuda"))
+pipe = StableDiffusionXLPipeline.from_pretrained(base, unet=unet, torch_dtype=torch.float16, variant="fp16").to("cuda")
+
+# Ensure sampler uses "trailing" timesteps.
+pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing")
+
+def generateTransparentImage(text):
+    # Ensure using the same inference steps as the loaded model and CFG set to 0.
+    image = pipe(text+', full body, transparent background', num_inference_steps=4, guidance_scale=0).images[0]
+    return image
+
+if __name__ == "__main__":
+    text = "a cat"
+    img = generateTransparentImage(text)
+    img.save("output.png")

requirements.txt

@@ -0,0 +1,12 @@
+omegaconf==2.3.0
+Pillow==10.1.0
+einops==0.7.0
+git+https://github.com/tatsy/torchmcubes.git
+diffusers["torch"]
+transformers==4.35.0
+trimesh==4.0.5
+rembg
+huggingface-hub
+imageio[ffmpeg]
+setuptools --upgrade
+torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

worker.py

@@ -0,0 +1,86 @@
+import subprocess
+from supabase import create_client, Client
+from imgGen import generateTransparentImage
+import sys
+sys.path.append('./TripoSR')
+
+import TripoSR.obj_gen as obj_gen
+import os
+from dotenv import load_dotenv
+import time
+load_dotenv()
+
+
+url: str = os.environ.get("SUPABASE_URL")
+key: str = os.environ.get("SUPABASE_KEY")
+supabase: Client = create_client(url, key)
+
+def check_queue():
+    try:
+        tasks = supabase.table("Tasks").select("*").eq("status", "pending").execute()
+        assert len(tasks.data) > 0
+        if len(tasks.data) > 0:
+            return tasks.data[0]
+        else:
+            return None
+    except Exception as e:
+        print(f"Error checking queue: {e}")
+        return None
+
+
+def generate_image(text):
+    try:
+        img = generateTransparentImage(text)
+        return img
+    except Exception as e:
+        print(f"Error generating image: {e}")
+        return None
+    
+
+def create_obj_file(img, task_id):
+    try:
+        obj_gen.generate_obj_from_image(img, 'task_'+str(task_id)+'.obj')
+    except Exception as e:
+        print(f"Error creating obj file: {e}")
+        supabase.table("Tasks").update({"status": "error"}).eq("id", task_id).execute()
+
+def send_back_to_supabase(task_id):
+    # check that a file was created
+    if os.path.exists('task_'+str(task_id)+'.obj'):
+        try:
+            with open('task_'+str(task_id)+'.obj', 'rb') as file:
+                data = file.read()
+                supabase.storage.from_('Results').upload('task_'+str(task_id)+'.obj', data)
+                public_url = supabase.storage.from_('Results').get_public_url('task_'+str(task_id)+'.obj')
+                supabase.table("Tasks").update({"status": "complete","result":public_url}).eq("id", task_id).execute()
+                os.remove('task_'+str(task_id)+'.obj')
+        except Exception as e:
+            print(f"Error sending file back to Supabase: {e}")
+            supabase.table("Tasks").update({"status": "error"}).eq("id", task_id).execute()
+
+    else:
+        print(f"Error: No file was created for task {task_id}")
+
+def worker():
+    while True:
+        task = check_queue()
+        if task:
+            supabase.table("Tasks").update({"status": "processing"}).eq("id", task['id']).execute()
+            print(f"Processing task {task['id']}")
+            img = generate_image(task["text"])
+            if img:
+                print(f"Image generated for task {task['id']}")
+                create_obj_file(img,task["id"])
+                send_back_to_supabase(task["id"])
+                print(f"Task {task['id']} completed")
+            else:
+                print(f"Error generating image for task {task['id']}")
+                supabase.table("Tasks").update({"status": "error"}).eq("id", task['id']).execute()
+
+        else:
+            print("No pending tasks in the queue")
+        
+        time.sleep(2)  # Add a 2 second delay between checks
+
+if __name__ == "__main__":
+    worker()