Update app.py

app.py

@@ -9,10 +9,11 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 
 try:
     print(f"Loading model on device: {device}")
-  
+    
     pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
         MODEL_ID, 
-        torch_dtype=torch.float16 if device == "cuda" else torch.float32
+        torch_dtype=torch.float16 if device == "cuda" else torch.float32,
+        low_cpu_mem_usage=True, # Added optimization for low RAM environments
     )
     pipe = pipe.to(device)
     print("Model loaded successfully.")
@@ -35,9 +36,7 @@ def generate_avatar(
     The 'style_strength' controls how much the new image deviates from the original.
     """
     if pipe is None:
-        return initial_image, "Error: The image generation model failed to load."
-
-    if not prompt:
+           if not prompt:
         return initial_image, "Please enter a descriptive text prompt."
         
     if initial_image is None:
@@ -46,7 +45,7 @@ def generate_avatar(
     print(f"Generating image with prompt: '{prompt}' and strength: {style_strength}")
     
     try:
-        with torch.autocast(device) if device == "cuda" else torch.no_grad():
+              with torch.autocast(device) if device == "cuda" else torch.no_grad():
             generated_image = pipe(
                 prompt=prompt,
                 image=initial_image.convert("RGB"), # Ensure image is in RGB format
@@ -62,6 +61,8 @@ def generate_avatar(
         print(f"Inference error: {e}")
         return initial_image, f"An error occurred during generation: {e}"
 
+
+# Define the input components
 image_input = gr.Image(type="pil", label="1. Upload Base Image (Your Photo or Style Reference)")
 prompt_input = gr.Textbox(
     label="2. Creative Prompt (e.g., 'A professional LinkedIn headshot, cyberpunk aesthetic, highly detailed')", 
@@ -88,7 +89,7 @@ negative_prompt_input = gr.Textbox(
 gr.Interface(
     fn=generate_avatar,
     inputs=[image_input, prompt_input, strength_slider, steps_slider, negative_prompt_input],
-    outputs=[gr.Image(type="pil", label="Generated Profile Picture "), gr.Textbox(label="Status")],
+    outputs=[gr.Image(type="pil", label="Generated Profile Picture"), gr.Textbox(label="Status")],
     title="✨ Creative Profile Picture Generator (SFW)",
     description="Upload a reference photo and a creative prompt to generate a new, stylized avatar using image-to-image diffusion. Transformation strength controls how much the output deviates from the input image (0.1 is subtle, 0.9 is highly stylized).",
     live=False,