Update pages/textimage.py

pages/textimage.py

@@ -1,69 +1,94 @@
 import streamlit as st
-import os
-import random
-import spaces
-import numpy as np
 import torch
-from typing import Tuple
-from datetime import datetime
-from diffusers import PixArtAlphaPipeline, LCMScheduler
+import numpy as np
+from PIL import Image
+import random
+import uuid
+from diffusers import PixArtAlphaPipeline
 
-# Check if CUDA is available
+# Check for CUDA availability
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
-# Define Hugging Face API details
-API_URL = "https://api-inference.huggingface.co/models/Huzaifa367/chat-summarizer"
-API_TOKEN = os.getenv("AUTH_TOKEN")
-HEADERS = {"Authorization": f"Bearer {API_TOKEN}"}
+# Load the PixArtAlphaPipeline
+if torch.cuda.is_available():
+    pipe = PixArtAlphaPipeline.from_pretrained(
+        "PixArt-alpha/PixArt-LCM-XL-2-1024-MS",
+        torch_dtype=torch.float16,
+        use_safetensors=True,
+    )
+    pipe.to(device)
+    st.write("Model loaded successfully!")
+else:
+    st.error("This demo requires GPU support, which is not available on this system.")
+
+# Constants
+MAX_SEED = np.iinfo(np.int32).max
+
+# Function to save image and return the path
+def save_image(img):
+    unique_name = str(uuid.uuid4()) + ".png"
+    img.save(unique_name)
+    return unique_name
+
+# Main function for image generation
+def generate_image(prompt, style, use_negative_prompt, negative_prompt, seed, width, height, inference_steps):
+    generator = torch.Generator().manual_seed(seed)
 
-# Initialize PixArtAlphaPipeline
-pipe = PixArtAlphaPipeline.from_pretrained(
-    "PixArt-alpha/PixArt-LCM-XL-2-1024-MS",
-    torch_dtype=torch.float16,
-    use_safetensors=True,
-    device=device
-)
+    # Apply the selected style
+    if style == "(No style)":
+        prompt_text = prompt
+    else:
+        prompt_text, _ = apply_style(style, prompt, negative_prompt)
 
-# Function to generate image based on prompt
-def generate_image(prompt: str) -> Tuple[str, int]:
-    seed = random.randint(0, np.iinfo(np.int32).max)
+    # Generate the image
     images = pipe(
-        prompt=prompt,
-        width=1024,
-        height=1024,
-        num_inference_steps=4,
-        generator=torch.Generator().manual_seed(seed),
+        prompt=prompt_text,
+        negative_prompt=None,
+        width=width,
+        height=height,
+        guidance_scale=0,
+        num_inference_steps=inference_steps,
+        generator=generator,
         num_images_per_prompt=1,
         use_resolution_binning=True,
         output_type="pil",
     ).images
 
-    # Save image and return path and seed
-    image_path = save_image(images[0])
-    return image_path, seed
+    # Save the image and display
+    if images:
+        img_path = save_image(images[0])
+        img = Image.open(img_path)
+        st.image(img, caption="Generated Image", use_column_width=True)
+        st.success("Image generated successfully!")
+    else:
+        st.error("Failed to generate image. Please try again.")
 
-# Function to save image and return path
-def save_image(img):
-    unique_name = str(uuid.uuid4()) + ".png"
-    img.save(unique_name)
-    return unique_name
+# Helper function to apply selected style
+def apply_style(style_name, positive, negative):
+    # Define styles dictionary (similar to your Gradio code)
+    styles = {
+        "(No style)": (positive, ""),
+        "Cinematic": ("cinematic still " + positive, "anime, cartoon, ..."),
+        "Realistic": ("Photorealistic " + positive, "drawing, painting, ..."),
+        # Add other styles here...
+    }
+    return styles.get(style_name, styles["(No style)"])
+
+# Streamlit UI
+st.title("Instant Image Generator")
 
-# Streamlit app
-def main():
-    st.set_page_config(layout="wide")
-    st.title("Instant Image Generator")
+prompt = st.text_input("Prompt", "Enter your prompt")
 
-    # Prompt input
-    prompt = st.text_area("Prompt", "Enter your prompt here...")
+style_names = ["(No style)", "Cinematic", "Realistic"]  # Add other styles here...
+style = st.selectbox("Image Style", style_names)
 
-    # Generate button
-    if st.button("Generate Image"):
-        if prompt:
-            # Generate image based on prompt
-            image_path, seed = generate_image(prompt)
+use_negative_prompt = st.checkbox("Use negative prompt")
+negative_prompt = st.text_input("Negative prompt", "")
 
-            # Display the generated image
-            st.image(image_path, use_column_width=True, caption=f"Seed: {seed}")
+seed = st.slider("Seed", 0, MAX_SEED, 0)
+width = st.slider("Width", 256, 4192, 1024, step=32)
+height = st.slider("Height", 256, 4192, 1024, step=32)
+inference_steps = st.slider("Steps", 4, 20, 4)
 
-if __name__ == "__main__":
-    main()
+if st.button("Generate Image"):
+    generate_image(prompt, style, use_negative_prompt, negative_prompt, seed, width, height, inference_steps)