Update app.py

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
 import spaces
 from gradio_litmodel3d import LitModel3D
+from diffusers import StableDiffusionPipeline
 
 import os
 os.environ['SPCONV_ALGO'] = 'native'
@@ -78,6 +79,19 @@ def unpack_state(state: dict) -> Tuple[Gaussian, edict, str]:
     return gs, mesh, state['trial_id']
 
 
+@spaces.GPU
+def text_to_image(prompt: str, seed: int, randomize_seed: bool) -> Image.Image:
+    """
+    Generate image from text prompt using Stable Diffusion.
+    """
+    if randomize_seed:
+        seed = np.random.randint(0, MAX_SEED)
+    
+    generator = torch.Generator(device="cuda").manual_seed(seed)
+    image = text2img_pipeline(prompt, generator=generator).images[0]
+    return image
+
+
 @spaces.GPU
 def image_to_3d(trial_id: str, seed: int, randomize_seed: bool, ss_guidance_strength: float, ss_sampling_steps: int, slat_guidance_strength: float, slat_sampling_steps: int) -> Tuple[dict, str]:
     """
@@ -160,6 +174,10 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
     
     with gr.Row():
         with gr.Column():
+            # Text to Image 부분 추가
+            text_prompt = gr.Textbox(label="Text Prompt", placeholder="Enter your text prompt here...")
+            generate_image_btn = gr.Button("Generate Image")
+            
             image_prompt = gr.Image(label="Image Prompt", image_mode="RGBA", type="pil", height=300)
             
             with gr.Accordion(label="Generation Settings", open=False):
@@ -205,6 +223,16 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
         )
 
     # Handlers
+    generate_image_btn.click(
+        text_to_image,
+        inputs=[text_prompt, seed, randomize_seed],
+        outputs=[image_prompt],
+    ).then(
+        preprocess_image,
+        inputs=[image_prompt],
+        outputs=[trial_id, image_prompt],
+    )
+
     image_prompt.upload(
         preprocess_image,
         inputs=[image_prompt],
@@ -248,6 +276,11 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
 if __name__ == "__main__":
     pipeline = TrellisImageTo3DPipeline.from_pretrained("JeffreyXiang/TRELLIS-image-large")
     pipeline.cuda()
+    
+    # Stable Diffusion pipeline 추가
+    text2img_pipeline = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
+    text2img_pipeline.to("cuda")
+    
     try:
         pipeline.preprocess_image(Image.fromarray(np.zeros((512, 512, 3), dtype=np.uint8)))    # Preload rembg
     except: