refactor: Update generate_image function to accept subject, style, color_scheme, angle, lighting_type, and additional_details as parameters

app.py

@@ -64,6 +64,28 @@ def preprocess(input_image):
     return input_image
 
 
+def generate_prompt(subject, style, color_scheme, angle, lighting_type, additional_details):
+    prompt = f"A 3D cartoon render of {subject}, featuring the entire body and shape, on a transparent background. The style should be {style}, with {color_scheme} colors, emphasizing the essential features and lines. The pose should clearly showcase the full form of the {subject} from a {angle} perspective. Lighting is {lighting_type}, highlighting the volume and depth of the subject. {additional_details}. Output as a high-resolution PNG with no background."
+    return prompt
+
+
+@spaces.GPU
+def generate_image(subject, style, color_scheme, angle, lighting_type, additional_details):
+    checkpoint = "sdxl_lightning_8step_unet.safetensors"
+    num_inference_steps = 8
+
+    pipe.scheduler = EulerDiscreteScheduler.from_config(
+        pipe.scheduler.config, timestep_spacing="trailing")
+    pipe.unet.load_state_dict(
+        load_file(hf_hub_download(repo, checkpoint), device="cuda"))
+
+    prompt = generate_prompt(subject, style, color_scheme,
+                             angle, lighting_type, additional_details)
+    results = pipe(
+        prompt, num_inference_steps=num_inference_steps, guidance_scale=0)
+    return results.images[0]
+
+
 @spaces.GPU
 def generate_mvs(input_image, sample_steps, sample_seed):
     seed_everything(sample_seed)
@@ -124,21 +146,6 @@ def make3d(images):
     return mesh_fpath, mesh_glb_fpath
 
 
-@spaces.GPU
-def generate_image(prompt):
-    checkpoint = "sdxl_lightning_8step_unet.safetensors"
-    num_inference_steps = 8
-
-    pipe.scheduler = EulerDiscreteScheduler.from_config(
-        pipe.scheduler.config, timestep_spacing="trailing")
-    pipe.unet.load_state_dict(
-        load_file(hf_hub_download(repo, checkpoint), device="cuda"))
-
-    results = pipe(
-        prompt, num_inference_steps=num_inference_steps, guidance_scale=0)
-    return results.images[0]
-
-
 # Configuration
 cuda_path = find_cuda()
 config_path = 'configs/instant-mesh-large.yaml'
@@ -193,8 +200,13 @@ with gr.Blocks() as demo:
     with gr.Group():
         with gr.Tab("Generate Image and Remove Background"):
             with gr.Row():
-                prompt = gr.Textbox(
-                    label='Enter your prompt (English)', scale=2)
+                subject = gr.Textbox(label='Subject', scale=2)
+                style = gr.Textbox(label='Style', scale=2)
+                color_scheme = gr.Textbox(label='Color Scheme', scale=2)
+                angle = gr.Textbox(label='Angle', scale=2)
+                lighting_type = gr.Textbox(label='Lighting Type', scale=2)
+                additional_details = gr.Textbox(
+                    label='Additional Details', scale=2)
                 submit_prompt = gr.Button(
                     'Generate Image', scale=1, variant='primary')
 
@@ -267,7 +279,7 @@ with gr.Blocks() as demo:
 
     mv_images = gr.State()
 
-    submit_prompt.click(fn=generate_image, inputs=[prompt], outputs=input_image).success(
+    submit_prompt.click(fn=generate_image, inputs=[subject, style, color_scheme, angle, lighting_type, additional_details], outputs=input_image).success(
         fn=preprocess, inputs=[input_image], outputs=[processed_image]
     )
     submit_process.click(fn=check_input_image, inputs=[input_image]).success(