Update app.py

app.py

@@ -48,29 +48,32 @@ def generate_text(prompt, temperature=0.9, max_new_tokens=512, top_p=0.95, repet
 def generate_image(prompt):
     try:
         result = image_client.predict(
-            prompt,  # Image prompt
-            "",  # Negative prompt
-            0,  # Seed
-            1024,  # Width
-            1024,  # Height
-            7.0,  # Guidance scale
-            28,  # Number of inference steps
-            'Euler a',  # Sampler
-            '896 x 1152',  # Aspect Ratio
-            '(None)',  # Style Preset
-            'Standard v3.1',  # Quality Tags Presets
-            False,  # Use Upscaler
-            0.55,  # Upscaler strength
-            1.5,  # Upscale by
-            True,  # Add Quality Tags
+            prompt=prompt,
+            negative_prompt="",
+            seed=0,
+            custom_width=1024,
+            custom_height=1024,
+            guidance_scale=7.0,
+            num_inference_steps=28,
+            sampler="Euler a",
+            aspect_ratio_selector="896 x 1152",
+            style_selector="(None)",
+            quality_selector="Standard v3.1",
+            use_upscaler=False,
+            upscaler_strength=0.55,
+            upscale_by=1.5,
+            add_quality_tags=True,
             isImg2Img=False,
-            img_path=None,  # Ensure img_path is set to None
+            img_path=None,
             img2img_strength=0.65,
             api_name="/run"
         )
-        # Check if result is not empty and contains the expected structure
-        if result and isinstance(result, list) and len(result) > 0 and 'image' in result[0]:
-            return result[0]['image']
+        # Process and display the result
+        if result:
+            images = []
+            for image_data in result[0]:
+                images.append(image_data['image'])
+            return images
         else:
             st.error("Unexpected result format from the Gradio API.")
             return None
@@ -107,8 +110,8 @@ def main():
         st.session_state.character_description = ""
     if "image_prompt" not in st.session_state:
         st.session_state.image_prompt = ""
-    if "image_path" not in st.session_state:
-        st.session_state.image_path = ""
+    if "image_paths" not in st.session_state:
+        st.session_state.image_paths = []
 
     # Generate button
     if st.button("Generate Waifu"):
@@ -123,7 +126,7 @@ def main():
             st.session_state.image_prompt = generate_text(image_prompt, temperature, max_new_tokens, top_p, repetition_penalty)
             
             # Generate image from image prompt
-            st.session_state.image_path = generate_image(st.session_state.image_prompt)
+            st.session_state.image_paths = generate_image(st.session_state.image_prompt)
             
             st.success("Waifu character generated!")
 
@@ -134,9 +137,10 @@ def main():
     if st.session_state.image_prompt:
         st.subheader("Image Prompt")
         st.write(st.session_state.image_prompt)
-    if st.session_state.image_path:
+    if st.session_state.image_paths:
         st.subheader("Generated Image")
-        st.image(st.session_state.image_path, caption="Generated Waifu Image")
+        for image_path in st.session_state.image_paths:
+            st.image(image_path, caption="Generated Waifu Image")
 
 if __name__ == "__main__":
     main()