Update app.py

app.py

@@ -0,0 +1,156 @@
+# Reference: https://huggingface.co/spaces/black-forest-labs/FLUX.1-schnell/blob/main/app.py
+import spaces
+import gradio as gr
+import numpy as np
+import random
+import torch
+import torch
+from diffusers import Transformer2DModel, PixArtSigmaPipeline, AutoencoderKL, DPMSolverMultistepScheduler, DDIMScheduler, EulerAncestralDiscreteScheduler, DPMSolverSDEScheduler
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, QuantoConfig, EetqConfig
+
+device = "cuda"
+weight_dtype = torch.float32
+weight_dtype_te = torch.bfloat16
+MAX_SEED = np.iinfo(np.int32).max
+
+vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=weight_dtype)
+scheduler=DPMSolverMultistepScheduler()
+pipe = PixArtSigmaPipeline(
+    vae=vae,
+    tokenizer=None,
+    text_encoder=None,
+    transformer=transformer,
+    scheduler=scheduler
+)
+
+pipe.to(device)
+
+tokenizer = AutoTokenizer.from_pretrained("cyberagent/calm2-7b")
+text_encoder =  AutoModelForCausalLM.from_pretrained(
+    "cyberagent/calm2-7b",
+    torch_dtype=weight_dtype_te,
+    device_map=device
+)
+
+@spaces.GPU()
+def infer(prompt, seed=42, randomize_seed=False, width=512, height=512, num_inference_steps=20, progress=gr.Progress(track_tqdm=True)):
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    generator = torch.Generator().manual_seed(seed)
+    with torch.no_grad():
+        pos_ids = tokenizer(
+            prompt, max_length=512, padding="max_length", truncation=True, return_tensors="pt",
+        ).to(device)
+        pos_emb = text_encoder(pos_ids.input_ids, output_hidden_states=True, attention_mask=pos_ids.attention_mask)
+        pos_emb = pos_emb.hidden_states[-1]
+        neg_ids = tokenizer(
+            "", max_length=512, padding="max_length", truncation=True, return_tensors="pt",
+        ).to(device)
+        neg_emb = text_encoder(neg_ids.input_ids, output_hidden_states=True, attention_mask=neg_ids.attention_mask)
+        neg_emb = neg_emb.hidden_states[-1]
+    
+        image = pipe(
+            negative_prompt=None,
+            prompt_embeds=pos_emb,
+            negative_prompt_embeds=neg_emb,
+            prompt_attention_mask=pos_ids.attention_mask,
+            negative_prompt_attention_mask=neg_ids.attention_mask,
+            max_sequence_length=512,
+            width=width,
+            height=height,
+            num_inference_steps=num_inference_steps,
+            generator=generator,
+            guidance_scale=4.5).images[0]
+    return image, seed
+ 
+examples = [
+    "芝生の上にあるピザ",
+    "東京の桜と建物。満開の桜の木が並び、ピンク色の花びらが風に舞っている。桜の背景には東京の高層ビルや伝統的な建物が調和して立っている。春の陽光が全体を明るく照らし、桜と建物が美しく映えている。都市の活気と自然の美しさが融合した風景。",
+    "パリは燃えているか",
+]
+
+css="""
+#col-container {
+    margin: 0 auto;
+    max-width: 520px;
+}
+"""
+
+with gr.Blocks(css=css) as demo:
+    
+    with gr.Column(elem_id="col-container"):
+        gr.Markdown(f"""# CommonArt β
+商用利用できる透明性の高い日本語画像生成AI
+        """)
+        
+        with gr.Row():
+            
+            prompt = gr.Text(
+                label="テキスト",
+                show_label=False,
+                max_lines=1,
+                placeholder="生成したいものを日本語や英語で説明してください",
+                container=False,
+            )
+            
+            run_button = gr.Button("生成", scale=0)
+        
+        result = gr.Image(label="生成結果", show_label=False)
+        
+        with gr.Accordion("詳細設定", open=False):
+            
+            seed = gr.Slider(
+                label="シード値",
+                minimum=0,
+                maximum=MAX_SEED,
+                step=1,
+                value=0,
+            )
+            
+            randomize_seed = gr.Checkbox(label="ランダム", value=True)
+            
+            with gr.Row():
+                
+                width = gr.Slider(
+                    label="幅",
+                    minimum=256,
+                    maximum=768,
+                    step=64,
+                    value=512,
+                )
+                
+                height = gr.Slider(
+                    label="高さ",
+                    minimum=256,
+                    maximum=768,
+                    step=64,
+                    value=512,
+                )
+            
+            with gr.Row():
+                
+  
+                num_inference_steps = gr.Slider(
+                    label="推論回数",
+                    minimum=1,
+                    maximum=50,
+                    step=1,
+                    value=20,
+                )
+        
+        gr.Examples(
+            examples = examples,
+            fn = infer,
+            inputs = [prompt],
+            outputs = [result, seed],
+            cache_examples="lazy"
+        )
+
+    gr.on(
+        triggers=[run_button.click, prompt.submit],
+        fn = infer,
+        inputs = [prompt, seed, randomize_seed, width, height, num_inference_steps],
+        outputs = [result, seed]
+    )
+
+demo.launch()