Update main.py

main.py

@@ -13,15 +13,10 @@ else:
     print(gpu_info)
     is_gpu = True
 print(is_gpu)
-
 from IPython.display import clear_output
-
-
-
 def check_enviroment():
     try:
         import torch
-
         print("Enviroment is already installed.")
     except ImportError:
         print("Enviroment not found. Installing...")
@@ -33,14 +28,9 @@ def check_enviroment():
         os.system("pip install python-dotenv")
         # Clear the output
         clear_output()
-
         print("Enviroment installed successfully.")
-
-
 # Call the function to check and install Packages if necessary
 check_enviroment()
-
-
 from IPython.display import clear_output
 import os
 import gradio as gr
@@ -49,16 +39,14 @@ import PIL
 import base64
 import io
 import torch
+import tempfile # Added for temporary file management
+
 # SDXL
 from diffusers import UNet2DConditionModel, DiffusionPipeline, LCMScheduler
 #requests
 import requests
 import random
 from PIL import Image
-
-
-
-
 # Get the current directory
 current_dir = os.getcwd()
 model_path = os.path.join(current_dir)
@@ -67,13 +55,10 @@ cache_path = os.path.join(current_dir, "cache")
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "1024"))
 SECRET_TOKEN = os.getenv("SECRET_TOKEN", "default_secret")
-
 API_TOKEN = os.environ.get("HF_READ_TOKEN")
 headers = {"Authorization": f"Bearer {API_TOKEN}"}
-
 # Uncomment the following line if you are using PyTorch 1.10 or later
 # os.environ["TORCH_USE_CUDA_DSA"] = "1"
-
 if is_gpu:
     # Uncomment the following line if you want to enable CUDA launch blocking
     os.environ["CUDA_LAUNCH_BLOCKING"] = "1"
@@ -85,7 +70,6 @@ current_dir = os.getcwd()
 model_path = os.path.join(current_dir)
 # Set the cache path
 cache_path = os.path.join(current_dir, "cache")
-
 def load_pipeline(use_cuda):
     device = "cuda" if use_cuda and torch.cuda.is_available() else "cpu"
     if device == "cuda":
@@ -99,12 +83,11 @@ def load_pipeline(use_cuda):
         pipe = DiffusionPipeline.from_pretrained("stabilityai/sdxl-turbo", use_safetensors=True)
         pipe = pipe.to(device)
     return pipe
-
 if is_sdxl:
     torch_dtype=torch.float16
     variant="fp16"
     unet = UNet2DConditionModel.from_pretrained(
-        "latent-consistency/lcm-sdxl",        
+        "latent-consistency/lcm-sdxl",
         torch_dtype=torch_dtype,
         variant=variant,
         cache_dir=cache_path,
@@ -135,12 +118,9 @@ if is_ssd:
     # load and fuse
     pipe.load_lora_weights("latent-consistency/lcm-lora-ssd-1b")
     pipe.fuse_lora()
-
 if is_sdxl_turbo:
     use_cuda=is_gpu
     pipe = load_pipeline(use_cuda)
-
-
 def generate(
     prompt: str,
     negative_prompt: str = "",
@@ -155,10 +135,7 @@ def generate(
         raise gr.Error(
             f"Invalid secret token. Please fork the original space if you want to use it for yourself."
         )
-
     generator = torch.Generator().manual_seed(seed)
-
-
     if not use_request:
         image = pipe(
             prompt=prompt,
@@ -170,7 +147,6 @@ def generate(
             generator=generator,
             output_type="pil",
         ).images[0]
-
     else:
         API_URL = "https://api-inference.huggingface.co/models/segmind/SSD-1B"
         payload = {
@@ -180,19 +156,19 @@ def generate(
             "cfg_scale": guidance_scale,
             "seed": seed if seed is not None else random.randint(-1, 2147483647)
         }
-
         image_bytes = requests.post(API_URL, headers=headers, json=payload).content
-        image = Image.open(io.BytesIO(image_bytes))    
-    return image
-
-
+        image = Image.open(io.BytesIO(image_bytes))
+        return image
 clear_output()
-
 from IPython.display import display
 
-
+# MODIFIED FUNCTION
 def generate_image(prompt="A beautiful and sexy girl",secret_token="default_secret"):
-    # Generate the image using the prompt
+    """
+    Generates an image, displays it, and immediately deletes the temporary file
+    to prevent storing images on disk.
+    """
+    # Generate the image in-memory using the prompt
     generated_image = generate(
         prompt=prompt,
         negative_prompt="",
@@ -203,9 +179,27 @@ def generate_image(prompt="A beautiful and sexy girl",secret_token="default_secr
         num_inference_steps=4,
         secret_token=secret_token
     )
-    # Display the image in the Jupyter Notebook
-    display(generated_image)
-
+    
+    # Create a temporary file to save the image.
+    # 'delete=False' allows us to manage its deletion manually.
+    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
+    temp_filepath = temp_file.name
+    
+    try:
+        # Save the generated image to the temporary file
+        generated_image.save(temp_filepath)
+        
+        # Display the image (this displays the in-memory object, not the file)
+        print("Displaying image...")
+        display(generated_image)
+        print("Image displayed.")
+
+    finally:
+        # This block ensures the file is always closed and deleted,
+        # even if errors occur.
+        temp_file.close()
+        os.remove(temp_filepath)
+        print(f"Temporary image file '{temp_filepath}' has been deleted.")
 
 if not run_api:
     secret_token = gr.Text(
@@ -228,7 +222,6 @@ if not run_api:
         visible=True,
     )
     seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
-
     width = gr.Slider(
         label="Width",
         minimum=256,
@@ -266,11 +259,8 @@ if not run_api:
         title="Image Generator",
         description="Generate images based on prompts.",
     )
-
     #iface.launch()
     iface.queue(max_size=32).launch(server_name="0.0.0.0", server_port=7860)  # Docker
-
-
 if run_api:
     with gr.Blocks() as demo:
         gr.HTML(
@@ -302,7 +292,6 @@ if run_api:
             visible=True,
         )
         seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
-
         width = gr.Slider(
             label="Width",
             minimum=256,
@@ -323,7 +312,6 @@ if run_api:
         num_inference_steps = gr.Slider(
             label="Number of inference steps", minimum=1, maximum=8, step=1, value=4
         )
-
         inputs = [
             prompt,
             negative_prompt,
@@ -340,16 +328,12 @@ if run_api:
             outputs=result,
             api_name="run",
         )
-
     # demo.queue(max_size=32).launch()
     # Launch the Gradio app with multiple workers and debug mode enabled
-    # demo.queue(max_size=32).launch(debug=True)# For Standard
+    # demo.queue(max_size=32).launch(debug=True)
+# For Standard
     demo.queue(max_size=32).launch(server_name="0.0.0.0", server_port=7860)  # Docker
-
-
 '''
-
-
 import gradio as gr
 import subprocess
 def run_command(command):
@@ -369,7 +353,7 @@ iface = gr.Interface(
     ["ls"],
     ["pwd"],
     ["echo 'Hello, Gradio!'"],
-    ["python --version"]]
-)
+    ["python --version"]
+])
 iface.launch(server_name="0.0.0.0", server_port=7860)
 '''