Update app.py

app.py

@@ -4,8 +4,6 @@ import random
 import spaces  # [uncomment to use ZeroGPU]
 from diffusers import DiffusionPipeline
 import torch
-import io
-from PIL import Image
 from tags import participant_tags, tribe_tags, skin_tone_tags, body_type_tags, tattoo_tags, piercing_tags, expression_tags, eye_tags, hair_style_tags, position_tags, fetish_tags, location_tags, camera_tags, atmosphere_tags
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -30,8 +28,10 @@ def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance
           selected_camera_tags, selected_atmosphere_tags, active_tab, progress=gr.Progress(track_tqdm=True)):
 
     if active_tab == "Prompt Input":
+        # Use the user-provided prompt
         final_prompt = f'score_9, score_8_up, score_7_up, source_anime, {prompt}'
     else:
+        # Use tags from the "Tag Selection" tab
         selected_tags = (
             [participant_tags[tag] for tag in selected_participant_tags] +
             [tribe_tags[tag] for tag in selected_tribe_tags] +
@@ -51,6 +51,7 @@ def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance
         tags_text = ', '.join(selected_tags)
         final_prompt = f'score_9, score_8_up, score_7_up, source_anime, {tags_text}'
 
+    # Concatenate user-provided negative prompt with additional restrictions
     additional_negatives = "worst quality, bad quality, jpeg artifacts, source_cartoon, 3d, (censor), monochrome, blurry, lowres, watermark"
     full_negative_prompt = f"{additional_negatives}, {negative_prompt}"
 
@@ -59,7 +60,7 @@ def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance
 
     generator = torch.Generator().manual_seed(seed)
 
-    # Generate the image
+    # Generate the image with the final prompts
     image = pipe(
         prompt=final_prompt,
         negative_prompt=full_negative_prompt,
@@ -70,14 +71,8 @@ def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance
         generator=generator
     ).images[0]
 
-    # Convert the image to the desired format (JPEG or PNG)
-    img_byte_arr = io.BytesIO()
-    image_format = "PNG"  # You can change this to "JPEG" if you prefer JPEG
-    image.save(img_byte_arr, format=image_format)
-    img_byte_arr.seek(0)  # Reset pointer to the beginning of the BytesIO object
-
-    # Return the image as a file (PNG or JPEG), seed, and the used prompts
-    return img_byte_arr, seed, f"Prompt used: {final_prompt}\nNegative prompt used: {full_negative_prompt}"
+    # Return image, seed, and the used prompts
+    return image, seed, f"Prompt used: {final_prompt}\nNegative prompt used: {full_negative_prompt}"
 
 
 css = """