Change architecture to SD2.1 and add predefined prompt lists

app.py

@@ -7,19 +7,23 @@ from diffusers import DiffusionPipeline
 import torch
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
-model_repo_id = "stabilityai/sdxl-turbo"
+model_repo_id = "stabilityai/stable-diffusion-2-1-base"
 
 torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
 
 pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
 pipe = pipe.to(device)
 
+backgrounds_list = ["forest", "city street", "beach", "office", "bus", "laboratory", "factory", "construction site", "hospital", "night club", ""]
+poses_list = ["portrait", "side-portrait"]
+
+gender_dict = {"ID_0": "male"}
 MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 1024
+MAX_IMAGE_SIZE = 512
 
 # @spaces.GPU  # Uncomment if using ZeroGPU
 def infer(
-    environment,
+    background,
     pose,
     negative_prompt,
     seed,
@@ -35,9 +39,12 @@ def infer(
 
     generator = torch.Generator().manual_seed(seed)
 
-    # Construct prompt from dropdown selections
-    prompt = f"A person {pose.lower()} in a {environment.lower()}, detailed, 8k"
+    id = "ID_0"
+    gender = gender_dict[id]
 
+    # Construct prompt from dropdown selections
+    prompt = f"face {pose.lower()} photo of {gender} {id} person, {background.lower()} background"
+    ""
     image = pipe(
         prompt=prompt,
         negative_prompt=negative_prompt,
@@ -59,19 +66,25 @@ css = """
 
 with gr.Blocks(css=css) as demo:
     with gr.Column(elem_id="col-container"):
-        gr.Markdown(" # Text-to-Image Gradio with Controlled Prompt")
+        gr.Markdown(" # ID-Booth Demo")
 
         with gr.Row():
-            environment = gr.Dropdown(
-                label="Environment",
-                choices=["Jungle", "Desert", "Space Station", "Underwater", "Urban City"],
+            which_id = gr.Dropdown(
+                label="Identity",
+                choices=backgrounds_list,
                 value="Jungle",
             )
 
+            background = gr.Dropdown(
+                label="Background",
+                choices=backgrounds_list,
+                value=backgrounds_list[0],
+            )
+
             pose = gr.Dropdown(
                 label="Pose",
-                choices=["Standing", "Sitting", "Running", "Flying", "Lying Down"],
-                value="Standing",
+                choices=poses_list,
+                value="Portrait",
             )
 
         run_button = gr.Button("Run", scale=0, variant="primary")
@@ -136,14 +149,14 @@ with gr.Blocks(css=css) as demo:
                 ["Space Station", "Flying"],
                 ["Urban City", "Sitting"],
             ],
-            inputs=[environment, pose],
+            inputs=[background, pose],
         )
 
     gr.on(
         triggers=[run_button.click],
         fn=infer,
         inputs=[
-            environment,
+            background,
             pose,
             negative_prompt,
             seed,