Update app.py

app.py

@@ -1,116 +1,116 @@
-import base64
-import re
-import os
-import random
-import requests
-import time
-from PIL import Image
-from io import BytesIO
-from typing import Tuple
-from diffusers import DiffusionPipeline
-import torch
-from datasets import load_dataset
-from dotenv import load_dotenv
-from flask import Flask, request, jsonify
-app = Flask(__name__)
-
-style_list = [
-    {
-        "name": "(No style)",
-        "prompt": "{prompt}",
-        "negative_prompt": "",
-    },
-    {
-        "name": "Cinematic",
-        "prompt": "cinematic still {prompt} . emotional, harmonious, vignette, highly detailed, high budget, bokeh, cinemascope, moody, epic, gorgeous, film grain, grainy",
-        "negative_prompt": "anime, cartoon, graphic, text, painting, crayon, graphite, abstract, glitch, deformed, mutated, ugly, disfigured",
-    },
-    {
-        "name": "Photographic",
-        "prompt": "cinematic photo {prompt} . 35mm photograph, film, bokeh, professional, 4k, highly detailed",
-        "negative_prompt": "drawing, painting, crayon, sketch, graphite, impressionist, noisy, blurry, soft, deformed, ugly",
-    },
-    {
-        "name": "Anime",
-        "prompt": "anime artwork {prompt} . anime style, key visual, vibrant, studio anime,  highly detailed",
-        "negative_prompt": "photo, deformed, black and white, realism, disfigured, low contrast",
-    },
-    {
-        "name": "Manga",
-        "prompt": "manga style {prompt} . vibrant, high-energy, detailed, iconic, Japanese comic style",
-        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast, realism, photorealistic, Western comic style",
-    },
-    {
-        "name": "Digital Art",
-        "prompt": "concept art {prompt} . digital artwork, illustrative, painterly, matte painting, highly detailed",
-        "negative_prompt": "photo, photorealistic, realism, ugly",
-    },
-    {
-        "name": "Pixel art",
-        "prompt": "pixel-art {prompt} . low-res, blocky, pixel art style, 8-bit graphics",
-        "negative_prompt": "sloppy, messy, blurry, noisy, highly detailed, ultra textured, photo, realistic",
-    },
-    {
-        "name": "Fantasy art",
-        "prompt": "ethereal fantasy concept art of  {prompt} . magnificent, celestial, ethereal, painterly, epic, majestic, magical, fantasy art, cover art, dreamy",
-        "negative_prompt": "photographic, realistic, realism, 35mm film, dslr, cropped, frame, text, deformed, glitch, noise, noisy, off-center, deformed, cross-eyed, closed eyes, bad anatomy, ugly, disfigured, sloppy, duplicate, mutated, black and white",
-    },
-    {
-        "name": "Neonpunk",
-        "prompt": "neonpunk style {prompt} . cyberpunk, vaporwave, neon, vibes, vibrant, stunningly beautiful, crisp, detailed, sleek, ultramodern, magenta highlights, dark purple shadows, high contrast, cinematic, ultra detailed, intricate, professional",
-        "negative_prompt": "painting, drawing, illustration, glitch, deformed, mutated, cross-eyed, ugly, disfigured",
-    },
-    {
-        "name": "3D Model",
-        "prompt": "professional 3d model {prompt} . octane render, highly detailed, volumetric, dramatic lighting",
-        "negative_prompt": "ugly, deformed, noisy, low poly, blurry, painting",
-    },
-]
-
-def infer(prompt, negative="low_quality", style_name=None, guidance_scale=None):
-
-
-    seed = random.randint(0,4294967295)
-    prompt, negative = apply_style(style_name, prompt, negative)
-    print(prompt)
-    print(negative)
-
-    # Load the Stable Diffusion model
-    pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, use_safetensors=True,variant="fp16")
-    pipe.to("cuda")
-
-    # Generate the images
-    images = pipe(prompt=prompt, negative_prompt=negative, guidance_scale=guidance_scale, seed=seed).images
-
-    # Convert the images to base64-encoded strings
-    image_urls = []
-    for i, image in enumerate(images):
-        buffered = BytesIO()
-        image.save(buffered, format="JPEG")
-        image_b64 = base64.b64encode(buffered.getvalue()).decode("utf-8")
-        #image_url = f"data:image/jpeg;base64,{image_b64}"
-        #image_urls.append(image_url)
-
-    return image_b64
-
-
-@app.route('/', methods=['POST'])
-def generate_image():
-    data = request.get_json()
-    if 'prompt' in data and 'style_name' in data and 'guidance_scale' in data:
-        prompt = data['prompt']
-        style_name = data['style_name']
-        guidance_scale = data['guidance_scale']
-        
-        image_urls = infer(prompt, style_name=style_name, guidance_scale=guidance_scale)
-        
-        # Convert the first generated image to base64
-        image_b64 = image_urls[0].split(",")[1]
-        
-        return jsonify({'image_base64': image_b64})
-    else:
-        return jsonify({'error': 'Missing required parameters'}), 400
-  
-
-if __name__ == '__main__':
-    app.run(host='0.0.0.0',port=6002,debug=True)
+import base64
+import re
+import os
+import random
+import requests
+import time
+from PIL import Image
+from io import BytesIO
+from typing import Tuple
+from diffusers import DiffusionPipeline
+import torch
+from datasets import load_dataset
+from dotenv import load_dotenv
+from flask import Flask, request, jsonify
+app = Flask(__name__)
+
+style_list = [
+    {
+        "name": "(No style)",
+        "prompt": "{prompt}",
+        "negative_prompt": "",
+    },
+    {
+        "name": "Cinematic",
+        "prompt": "cinematic still {prompt} . emotional, harmonious, vignette, highly detailed, high budget, bokeh, cinemascope, moody, epic, gorgeous, film grain, grainy",
+        "negative_prompt": "anime, cartoon, graphic, text, painting, crayon, graphite, abstract, glitch, deformed, mutated, ugly, disfigured",
+    },
+    {
+        "name": "Photographic",
+        "prompt": "cinematic photo {prompt} . 35mm photograph, film, bokeh, professional, 4k, highly detailed",
+        "negative_prompt": "drawing, painting, crayon, sketch, graphite, impressionist, noisy, blurry, soft, deformed, ugly",
+    },
+    {
+        "name": "Anime",
+        "prompt": "anime artwork {prompt} . anime style, key visual, vibrant, studio anime,  highly detailed",
+        "negative_prompt": "photo, deformed, black and white, realism, disfigured, low contrast",
+    },
+    {
+        "name": "Manga",
+        "prompt": "manga style {prompt} . vibrant, high-energy, detailed, iconic, Japanese comic style",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast, realism, photorealistic, Western comic style",
+    },
+    {
+        "name": "Digital Art",
+        "prompt": "concept art {prompt} . digital artwork, illustrative, painterly, matte painting, highly detailed",
+        "negative_prompt": "photo, photorealistic, realism, ugly",
+    },
+    {
+        "name": "Pixel art",
+        "prompt": "pixel-art {prompt} . low-res, blocky, pixel art style, 8-bit graphics",
+        "negative_prompt": "sloppy, messy, blurry, noisy, highly detailed, ultra textured, photo, realistic",
+    },
+    {
+        "name": "Fantasy art",
+        "prompt": "ethereal fantasy concept art of  {prompt} . magnificent, celestial, ethereal, painterly, epic, majestic, magical, fantasy art, cover art, dreamy",
+        "negative_prompt": "photographic, realistic, realism, 35mm film, dslr, cropped, frame, text, deformed, glitch, noise, noisy, off-center, deformed, cross-eyed, closed eyes, bad anatomy, ugly, disfigured, sloppy, duplicate, mutated, black and white",
+    },
+    {
+        "name": "Neonpunk",
+        "prompt": "neonpunk style {prompt} . cyberpunk, vaporwave, neon, vibes, vibrant, stunningly beautiful, crisp, detailed, sleek, ultramodern, magenta highlights, dark purple shadows, high contrast, cinematic, ultra detailed, intricate, professional",
+        "negative_prompt": "painting, drawing, illustration, glitch, deformed, mutated, cross-eyed, ugly, disfigured",
+    },
+    {
+        "name": "3D Model",
+        "prompt": "professional 3d model {prompt} . octane render, highly detailed, volumetric, dramatic lighting",
+        "negative_prompt": "ugly, deformed, noisy, low poly, blurry, painting",
+    },
+]
+
+def infer(prompt, negative="low_quality", style_name=None, guidance_scale=None):
+
+
+    seed = random.randint(0,4294967295)
+    prompt, negative = apply_style(style_name, prompt, negative)
+    print(prompt)
+    print(negative)
+
+    # Load the Stable Diffusion model
+    pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, use_safetensors=True,variant="fp16")
+    pipe.to("cuda")
+
+    # Generate the images
+    images = pipe(prompt=prompt, negative_prompt=negative, guidance_scale=guidance_scale, seed=seed).images
+
+    # Convert the images to base64-encoded strings
+    image_urls = []
+    for i, image in enumerate(images):
+        buffered = BytesIO()
+        image.save(buffered, format="JPEG")
+        image_b64 = base64.b64encode(buffered.getvalue()).decode("utf-8")
+        #image_url = f"data:image/jpeg;base64,{image_b64}"
+        #image_urls.append(image_url)
+
+    return image_b64
+
+
+@app.route('/', methods=['POST'])
+def generate_image():
+    data = request.get_json()
+    if 'prompt' in data and 'style_name' in data and 'guidance_scale' in data:
+        prompt = data['prompt']
+        style_name = data['style_name']
+        guidance_scale = data['guidance_scale']
+        
+        image_urls = infer(prompt, style_name=style_name, guidance_scale=guidance_scale)
+        
+        # Convert the first generated image to base64
+        image_b64 = image_urls[0].split(",")[1]
+        
+        return jsonify({'image_base64': image_b64})
+    else:
+        return jsonify({'error': 'Missing required parameters'}), 400
+  
+
+if __name__ == '__main__':
+    app.run(host='0.0.0.0',port=7860,debug=True)